More bugfixes for statespace

jessegrabowski commented 1 month ago

The statespace model is still broken, this PR is another round of bugfixes.

[x] Update structural notebook to conform to API changes after #326
[x] Remove the numpy helper functions from the structural notebook, use pytensor only
[x] Fix JAX-based forward sampling
[x] Re-run all notebooks without errors
[x] Get all current tests to pass
[x] Add JAX tests

I need some help with fixing the JAX forward sampling. I'm doing something wrong, because even after freezing I have dynamic shape errors. This is the major blocker to considering the module "working" again.

review-notebook-app[bot] commented 1 month ago

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

jessegrabowski commented 1 week ago

This is ready for review. I will need to make changes again after the next pymc/pytensor releases, but for now this all works -- tests pass, and all notebooks run.

For review I tried to organize the commits into chunks. It would be good to have eyeballs on the changes related to the distributions, since those were creating trouble for me in the first place.

jessegrabowski commented 1 week ago

No idea what's going on with the CI, there seems to be a lot of broken stuff unrelated to this.

pymc-devs / pymc-experimental

More bugfixes for statespace #346