maintenance of statespace notebooks

[Dekermanjian]

We have had a few rounds of making a bit of changes to the statespace module. Adding the multivariate support to structural, updating naming conventions, and changing how we specify different backends to be used. These changes are not reflected in the example notebooks and in certain cases introduce errors in certain cells. These notebooks are possibly the primary means in which new users learn how to use the state space module and would be beneficial to keep them maintained.

[jessegrabowski]

100% true. Also add to this the Hurricane example in pymc-examples, which is our highest visibility use of statespace.

[Dekermanjian]

Yeah, I forgot about the hurricane! I will do a pass over it today and update it over on pymc-examples.

[theorashid]

Hey all,

I've been trying to get the joint estimation ETS notebook section running today. I am using pytensor 2.32.0 so that I have Jesse's BlockDiagonal rewrites. Just running this though, I am getting some pytensor and then some jax errors. I know @AlexAndorra said he had some pytensor errors so wondering whether these are expected or the same (hate to do "it doesn't work on my machine", but posting in case it was a shared experience).

Note, this might also be to do with my jax deps. I have an annoying system constraint of GCC 7.3.1 which has been causing some headaches. Here are the errors:

Building PyMC model and sampling...
/XXXXXX/src/pymc_ss/.venv/lib/python3.12/site-packages/pymc_extras/statespace/utils/data_tools.py:175: ImputationWarning: Provided data contains missing values and will be automatically imputed as hidden states during Kalman filtering.
  warnings.warn(impute_message, ImputationWarning)
Sampling from posterior...
/XXXXXX/src/pymc_ss/.venv/lib/python3.12/site-packages/pytensor/link/c/cmodule.py:2968: UserWarning: PyTensor could not link to a BLAS installation. Operations that might benefit from BLAS will be severely degraded.
This usually happens when PyTensor is installed via pip. We recommend it be installed via conda/mamba/pixi instead.
Alternatively, you can use an experimental backend such as Numba or JAX that perform their own BLAS optimizations, by setting `pytensor.config.mode == 'NUMBA'` or passing `mode='NUMBA'` when compiling a PyTensor function.
For more options and details see https://pytensor.readthedocs.io/en/latest/troubleshooting.html#how-do-i-configure-test-my-blas-library
  warnings.warn(
<<!! BUG IN FGRAPH.REPLACE OR A LISTENER !!>> <class 'TypeError'> The type of the replacement (Matrix(float64, shape=(6, 2))) must be compatible with the type of the original Variable (Matrix(float64, shape=(6, 6))). local_block_diag_dot_to_dot_block_diag
ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_block_diag_dot_to_dot_block_diag
ERROR (pytensor.graph.rewriting.basic): node: Blockwise{BlockDiagonal{n_inputs=2}, (m0,n0),(m1,n1)->(m,n)}(SetSubtensor{i, j}.0, SetSubtensor{i, j}.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/XXXXXX/src/pymc_ss/.venv/lib/python3.12/site-packages/pytensor/graph/rewriting/basic.py", line 1979, in process_node
    fgraph.replace_all_validate_remove(  # type: ignore
  File "/XXXXXX/src/pymc_ss/.venv/lib/python3.12/site-packages/pytensor/graph/features.py", line 800, in replace_all_validate_remove
    chk = fgraph.replace_all_validate(replacements, reason=reason, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/XXXXXX/src/pymc_ss/.venv/lib/python3.12/site-packages/pytensor/graph/features.py", line 745, in replace_all_validate
    fgraph.replace(r, new_r, reason=reason, verbose=False, **kwargs)
  File "/XXXXXX/src/pymc_ss/.venv/lib/python3.12/site-packages/pytensor/graph/fg.py", line 535, in replace
    self.change_node_input(
  File "/XXXXXX/src/pymc_ss/.venv/lib/python3.12/site-packages/pytensor/graph/fg.py", line 447, in change_node_input
    raise TypeError(
TypeError: The type of the replacement (Matrix(float64, shape=(6, 2))) must be compatible with the type of the original Variable (Matrix(float64, shape=(6, 6))).

...
ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_block_diag_dot_to_dot_block_diag
ERROR (pytensor.graph.rewriting.basic): node: Blockwise{BlockDiagonal{n_inputs=2}, (m0,n0),(m1,n1)->(m,n)}(SetSubtensor{i, j}.0, SetSubtensor{i, j}.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):

But hey, it does sample (the line idata_joint = pm.sample(**sampler_kwargs())) after that. But then we get the jax erors when I try and forecast idata_joint_post = ss_mod_joint.sample_conditional_posterior(idata_joint); idata_joint_forecast = ss_mod_joint.forecast(idata_joint, start=df_train.index[-1], end=df_test.index[-1])

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_subtensor_of_squeeze
ERROR (pytensor.graph.rewriting.basic): node: Subtensor{:stop}(Squeeze{axis=2}.0, ScalarFromTensor.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/XXXXXX/src/pymc_ss/.venv/lib/python3.12/site-packages/pytensor/graph/rewriting/basic.py", line 1933, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/XXXXXX/src/pymc_ss/.venv/lib/python3.12/site-packages/pytensor/graph/rewriting/basic.py", line 1086, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/XXXXXX/src/pymc_ss/.venv/lib/python3.12/site-packages/pytensor/tensor/rewriting/subtensor_lift.py", line 491, in local_subtensor_of_squeeze
    new_idxs = np.insert(np.array(idxs, dtype=object), dropped_dims, slice(None))
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/XXXXXX/src/pymc_ss/.venv/lib/python3.12/site-packages/numpy/lib/_function_base_impl.py", line 5589, in insert
    raise IndexError(f"index {obj} is out of bounds for axis {axis} "
IndexError: index [2] is out of bounds for axis 0 with size 1

<<!! BUG IN FGRAPH.REPLACE OR A LISTENER !!>> <class 'TypeError'> The type of the replacement (Matrix(float64, shape=(6, 2))) must be compatible with the type of the original Variable (Matrix(float64, shape=(6, 6))). local_block_diag_dot_to_dot_block_diag
ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_block_diag_dot_to_dot_block_diag
ERROR (pytensor.graph.rewriting.basic): node: Blockwise{BlockDiagonal{n_inputs=2}, (m0,n0),(m1,n1)->(m,n)}(SetSubtensor{i, j}.0, SetSubtensor{i, j}.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/XXXXXX/src/pymc_ss/.venv/lib/python3.12/site-packages/pytensor/graph/rewriting/basic.py", line 1979, in process_node
    fgraph.replace_all_validate_remove(  # type: ignore
  File "/XXXXXX/src/pymc_ss/.venv/lib/python3.12/site-packages/pytensor/graph/features.py", line 800, in replace_all_validate_remove
    chk = fgraph.replace_all_validate(replacements, reason=reason, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/XXXXXX/src/pymc_ss/.venv/lib/python3.12/site-packages/pytensor/graph/features.py", line 745, in replace_all_validate
    fgraph.replace(r, new_r, reason=reason, verbose=False, **kwargs)
  File "/XXXXXX/src/pymc_ss/.venv/lib/python3.12/site-packages/pytensor/graph/fg.py", line 535, in replace
    self.change_node_input(
  File "/XXXXXX/src/pymc_ss/.venv/lib/python3.12/site-packages/pytensor/graph/fg.py", line 447, in change_node_input
    raise TypeError(
TypeError: The type of the replacement (Matrix(float64, shape=(6, 2))) must be compatible with the type of the original Variable (Matrix(float64, shape=(6, 6))).
...
Sampling: [filtered_posterior, filtered_posterior_observed, predicted_posterior, predicted_posterior_observed, smoothed_posterior, smoothed_posterior_observed]
Sampling ... ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━   0% -:--:-- / 0:00:00
Traceback (most recent call last):
  File "/XXXXXX/src/pymc_ss/.venv/lib/python3.12/site-packages/pytensor/link/basic.py", line 665, in thunk
    outputs = fgraph_jit(*(x[0] for x in thunk_inputs))
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/tmp/tmpgjua90ko", line 393, in jax_funcified_fgraph
    tensor_variable_359, tensor_variable_360, tensor_variable_361, tensor_variable_362, tensor_variable_363, tensor_variable_364, tensor_variable_365, tensor_variable_366, tensor_variable_367, tensor_variable_368, tensor_variable_369, tensor_variable_370, filtered_posterior_observed, tensor_variable_371, tensor_variable_372, tensor_variable_373, tensor_variable_374, variable_1 = scan(tensor_constant_43, tensor_constant_44, tensor_constant_45, tensor_constant_46, tensor_variable_179, tensor_variable_176, tensor_variable_168, tensor_variable_159, tensor_variable_152, tensor_variable_144, tensor_variable_124, tensor_variable_115, tensor_variable_94, tensor_variable_85, tensor_variable_64, tensor_variable_51, random_generator_shared_variable, tensor_constant_43, tensor_constant_43, tensor_constant_43, tensor_constant_43, tensor_constant_43, tensor_constant_43, tensor_constant_43, transition, tensor_variable_27, tensor_variable_4)
                                                                                                                                                                                                                                                                                                                                                                                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/XXXXXX/src/pymc_ss/.venv/lib/python3.12/site-packages/pytensor/link/jax/dispatch/scan.py", line 141, in scan
    final_carry, traces = jax.lax.scan(
                          ^^^^^^^^^^^^^
ValueError: scan got `length` argument of 182 which disagrees with leading axis sizes [182, 182, 182, 6, 1].
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/XXXXXX/src/pymc_ss/joint_estimation_example.py", line 218, in <module>
    if __name__ == "__main__":
                  ^^^^^^^^^^^^^
  File "/XXXXXX/src/pymc_ss/joint_estimation_example.py", line 170, in run_joint_estimation_example
    # Generate posterior predictions and forecasts
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/XXXXXX/src/pymc_ss/.venv/lib/python3.12/site-packages/pymc_extras/statespace/core/statespace.py", line 1476, in sample_conditional_posterior
    return self._sample_conditional(
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/XXXXXX/src/pymc_ss/.venv/lib/python3.12/site-packages/pymc_extras/statespace/core/statespace.py", line 1251, in _sample_conditional
    idata_conditional = pm.sample_posterior_predictive(
                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/XXXXXX/src/pymc_ss/.venv/lib/python3.12/site-packages/pymc/sampling/forward.py", line 951, in sample_posterior_predictive
    values = sampler_fn(**param)
             ^^^^^^^^^^^^^^^^^^^
  File "/XXXXXX/src/pymc_ss/.venv/lib/python3.12/site-packages/pymc/util.py", line 390, in wrapped
    return core_function(**input_point)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/XXXXXX/src/pymc_ss/.venv/lib/python3.12/site-packages/pytensor/compile/function/types.py", line 1039, in __call__
    outputs = vm() if output_subset is None else vm(output_subset=output_subset)
              ^^^^
  File "/XXXXXX/src/pymc_ss/.venv/lib/python3.12/site-packages/pytensor/link/basic.py", line 669, in thunk
    raise_with_op(self.fgraph, output_nodes[0], thunk)
  File "/XXXXXX/src/pymc_ss/.venv/lib/python3.12/site-packages/pytensor/link/utils.py", line 526, in raise_with_op
    raise exc_value.with_traceback(exc_trace)
  File "/XXXXXX/src/pymc_ss/.venv/lib/python3.12/site-packages/pytensor/link/basic.py", line 665, in thunk
    outputs = fgraph_jit(*(x[0] for x in thunk_inputs))
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/tmp/tmpgjua90ko", line 393, in jax_funcified_fgraph
    tensor_variable_359, tensor_variable_360, tensor_variable_361, tensor_variable_362, tensor_variable_363, tensor_variable_364, tensor_variable_365, tensor_variable_366, tensor_variable_367, tensor_variable_368, tensor_variable_369, tensor_variable_370, filtered_posterior_observed, tensor_variable_371, tensor_variable_372, tensor_variable_373, tensor_variable_374, variable_1 = scan(tensor_constant_43, tensor_constant_44, tensor_constant_45, tensor_constant_46, tensor_variable_179, tensor_variable_176, tensor_variable_168, tensor_variable_159, tensor_variable_152, tensor_variable_144, tensor_variable_124, tensor_variable_115, tensor_variable_94, tensor_variable_85, tensor_variable_64, tensor_variable_51, random_generator_shared_variable, tensor_constant_43, tensor_constant_43, tensor_constant_43, tensor_constant_43, tensor_constant_43, tensor_constant_43, tensor_constant_43, transition, tensor_variable_27, tensor_variable_4)
                                                                                                                                                                                                                                                                                                                                                                                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/XXXXXX/src/pymc_ss/.venv/lib/python3.12/site-packages/pytensor/link/jax/dispatch/scan.py", line 141, in scan
    final_carry, traces = jax.lax.scan(
                          ^^^^^^^^^^^^^
ValueError: scan got `length` argument of 182 which disagrees with leading axis sizes [182, 182, 182, 6, 1].
Apply node that caused the error: Scan{scan_fn&scan_fn&scan_fn&scan_fn&scan_fn, while_loop=False, inplace=none}(182, Scan{forward_kalman_pass&scan_fn&forward_kalman_pass&forward_kalman_pass&forward_kalman_pass&forward_kalman_pass, while_loop=False, inplace=none}.10, Scan{forward_kalman_pass&scan_fn&forward_kalman_pass&forward_kalman_pass&forward_kalman_pass&forward_kalman_pass, while_loop=False, inplace=none}.11, predicted_states, Subtensor{:stop}.0, Dot.0, Reshape{3}.0, smoothed_states, Subtensor{:stop}.0, Dot.0, Reshape{3}.0, RNG(<Generator(PCG64) at 0x7F872E6C4E40>), RNG(<Generator(PCG64) at 0x7F872E6C4BA0>), RNG(<Generator(PCG64) at 0x7F872E6C4660>), RNG(<Generator(PCG64) at 0x7F872E6C4580>), RNG(<Generator(PCG64) at 0x7F872E6C44A0>), 182, 182, 182, 182, 182)
Toposort index: 256
Inputs types: [TensorType(int64, shape=()), TensorType(float64, shape=(None, 6)), TensorType(float64, shape=(None, 6, 6)), TensorType(float64, shape=(None, 6)), TensorType(float64, shape=(None, 6, 6)), TensorType(float64, shape=(None, 2)), TensorType(float64, shape=(None, 2, 2)), TensorType(float64, shape=(None, 6)), TensorType(float64, shape=(None, 6, 6)), TensorType(float64, shape=(None, 2)), TensorType(float64, shape=(None, 2, 2)), RandomGeneratorType, RandomGeneratorType, RandomGeneratorType, RandomGeneratorType, RandomGeneratorType, TensorType(int64, shape=()), TensorType(int64, shape=()), TensorType(int64, shape=()), TensorType(int64, shape=()), TensorType(int64, shape=())]
Inputs shapes: [(2, 2), (2,), (2,), (2,), (2,), (2,), 'No shapes', 'No shapes', 'No shapes', 'No shapes', 'No shapes', 'No shapes']
Inputs strides: [(16, 8), (8,), (8,), (8,), (8,), (8,), 'No strides', 'No strides', 'No strides', 'No strides', 'No strides', 'No strides']
Inputs values: [array([[24.95366585, -0.38850447],
       [-0.38850447,  0.06003416]]), array([0.9909915 , 0.88264888]), array([0.73555206, 0.51269118]), array([0.23003079, 0.0231709 ]), array([-12.25669939,  -0.14883723]), array([26.54918196, 20.31707694]), {'bit_generator': 1, 'state': {'state': np.uint32(1503019118), 'inc': np.uint32(3404321205)}, 'has_uint32': 0, 'uinteger': 0, 'jax_state': array([ 973644868, 1503019118], dtype=uint32)}, {'bit_generator': 1, 'state': {'state': np.uint32(2385338737), 'inc': np.uint32(735519806)}, 'has_uint32': 0, 'uinteger': 0, 'jax_state': array([4013715060, 2385338737], dtype=uint32)}, {'bit_generator': 1, 'state': {'state': np.uint32(1776780982), 'inc': np.uint32(1315367445)}, 'has_uint32': 0, 'uinteger': 0, 'jax_state': array([2934872578, 1776780982], dtype=uint32)}, {'bit_generator': 1, 'state': {'state': np.uint32(3942769424), 'inc': np.uint32(4087074669)}, 'has_uint32': 0, 'uinteger': 0, 'jax_state': array([3329923146, 3942769424], dtype=uint32)}, {'bit_generator': 1, 'state': {'state': np.uint32(957240127), 'inc': np.uint32(2747591359)}, 'has_uint32': 0, 'uinteger': 0, 'jax_state': array([2606351377,  957240127], dtype=uint32)}, {'bit_generator': 1, 'state': {'state': np.uint32(99726969), 'inc': np.uint32(3087850819)}, 'has_uint32': 0, 'uinteger': 0, 'jax_state': array([3349574534,   99726969], dtype=uint32)}]
Outputs clients: [[output[0](filtered_posterior)], [output[2](predicted_posterior)], [output[3](predicted_posterior_observed)], [output[4](smoothed_posterior)], [output[5](smoothed_posterior_observed)], [output[6](Scan{scan_fn&scan_fn&scan_fn&scan_fn&scan_fn, while_loop=False, inplace=none}.5)], [output[8](Scan{scan_fn&scan_fn&scan_fn&scan_fn&scan_fn, while_loop=False, inplace=none}.6)], [output[9](Scan{scan_fn&scan_fn&scan_fn&scan_fn&scan_fn, while_loop=False, inplace=none}.7)], [output[10](Scan{scan_fn&scan_fn&scan_fn&scan_fn&scan_fn, while_loop=False, inplace=none}.8)], [output[11](Scan{scan_fn&scan_fn&scan_fn&scan_fn&scan_fn, while_loop=False, inplace=none}.9)]]

HINT: Re-running with most PyTensor optimizations disabled could provide a back-trace showing when this node was created. This can be done by setting the PyTensor flag 'optimizer=fast_compile'. If that does not work, PyTensor optimizations can be disabled with 'optimizer=None'.
HINT: Use the PyTensor flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.

Here's a summary of my env

CORE BAYESIAN STACK:
  PyMC:           5.25.1+9.g09603234f
  PyTensor:       2.32.0 ✅ (with block diagonal optimizations)
  ArviZ:          0.22.0
  Preliz:         0.20.0

HIGH-PERFORMANCE COMPUTING:
  Nutpie:         0.15.2 ✅ (NUTS sampler)
  JAX:            0.4.38 ✅ (backend)
  JAXlib:         0.4.38 ✅ (JAX library)

SYSTEM COMPATIBILITY:
  NumPy:          2.2.1 ✅ (GCC 7.3.1 compatible)
  Python:         3.12
  GCC:            7.3.1 (compatible)

[Dekermanjian]

Hey @theorashid, I am checking this now. I am on pymc 5.25.1 and pytensor 2.31.7. I am going to try to run through the ETS notebook and then upgrade pytensor and run through again. Give me a few minutes and I will let you know if I am also seeing the same error on my end.

[theorashid]

Also, a standard question that I've never understood. We have jax flags everywhere including with the nutpie sampler. Why do I still get this warning if I'm saying not to use the C backend

src/pymc_ss/.venv/lib/python3.12/site-packages/pytensor/link/c/cmodule.py:2968: UserWarning: PyTensor could not link to a BLAS installation. Operations that might benefit from BLAS will be severely degraded.
This usually happens when PyTensor is installed via pip. We recommend it be installed via conda/mamba/pixi instead.
Alternatively, you can use an experimental backend such as Numba or JAX that perform their own BLAS optimizations, by setting `pytensor.config.mode == 'NUMBA'` or passing `mode='NUMBA'` when compiling a PyTensor function.
For more options and details see https://pytensor.readthedocs.io/en/latest/troubleshooting.html#how-do-i-configure-test-my-blas-library

[Dekermanjian]

Hey @theorashid, I was able to run the ETS notebook with no issues using pymc 5.25.1 and pytensor 2.31.7. I am trying to install pytensor 2.32.0 but it just isn't wanting to install. Did you use coda-forge to install it?

[AlexAndorra]

Yeah, that's the same error I flagged the other day @theorashid . @ricardoV94 sent me the line I should add a dprint to in the pytensor code base, but haven't gotten the time to do it yet. Will be soon, but feel free to beat me to it if you have bandwidth now! Ping me on Discord if you don't find the public thread.

@Dekermanjian install the main pytensor branch, and you will be able to reproduce the error

[ricardoV94]

We have jax flags everywhere including with the nutpie sampler. Why do I still get this warning if I'm saying not to use the C backend

You're still compiling some stuff to C, perhaps it's prior or posterior predictive sampling? what function call triggers the warning?

[Dekermanjian]

install the main pytensor branch, and you will be able to reproduce the error

Thank you, Alex. I will give that a shot today. I am interested to see the speed up in multivariate use cases.

[Dekermanjian]

@jessegrabowski I did a pass through on the hurricanes notebook. The only thing that I saw that would need updating is the way in which we are pulling the KF outputs from the model. We now have the new sample_filter_outputs() helper method but that isn't released yet. As soon as that gets released I will push and update to the notebook with that change.