Statespace: Don't automatically save statespace matrices as Deterministic variables

[jessegrabowski]

I had originally done this to facilitate out-of-sample sampling tasks, so I could do e.g. pm.Flat('T', T). The result was that all matrices were saved to the idata, creating pretty horrible pm.model_to_graphviz outputs like this:

This also was extremely memory wasteful. Many of the matrices are not random at all, and they were being saved (chain, draw) times.

After this refactor, the matrices are dynamically rebuilt as needed from the parameter samples. The new graphs look like this:

There are also some en-passant changes to how exogenous data are handled that breaks the Structural example notebook. I will open a new PR after this one to address that, because I think I finally have it set up to handle forecasting with exogenous data. Basically, I was previously treating exogenous data like a type of "parameter". It's been upgraded to a first-class object, and custom models subclassing PyMCStateSpace that use exogenous data will now need to implement data_names and data_info properties.

[ricardoV94]

For the memory question, matrices that don't change could be defined as constant or mutabledata

[jessegrabowski]

They're stored as TensorVariables in the statespace representation. This new way of doing this will directly use those once instead of copying them over and over into the idata.

[ricardoV94]

They're stored as TensorVariables in the statespace representation. This new way of doing this will directly use those once instead of copying them over and over into the idata.

ConstantData and MutableData are also stored only once in the idata as opposed to Deterministics which are stored per draw. Seems like that's what you wanted?

[jessegrabowski]

Something like that, but I don't want to have a bunch of logic to decide if a matrix is static or contains parameters. If the user really wants to inspect matrices, he can ask for them manually and save them however he wants. In general, I think the important outputs for most users is going to be the parameters and the states. The rest can be more hidden away.

[ricardoV94]

Just a tiny type-hint thing

[ricardoV94]

List doesn't need ellipsis

Suggested change

	def unpack_statespace(self) -> list[pt.TensorVariable, ...]:
	def unpack_statespace(self) -> list[pt.TensorVariable]:

[jessegrabowski]

Pycharm actually complains when I make this change, it wants me to literally list out tuple[ TensorVariable, TensorVariable, TensorVariable, TensorVariable, TensorVariable, TensorVariable, TensorVariable, TensorVariable, TensorVariable], which seems silly. list[TensorVariable, ...] was my hack-around.

[ricardoV94]

Ignore pycharm

[jessegrabowski]

What's up with these jax test failures on the Ubuntu CI?

[ricardoV94]

What's up with these jax test failures on the Ubuntu CI?

#305 ?