search

CDCgov / PyRenew

Python package for multi-signal Bayesian renewal modeling with JAX and NumPyro.
https://cdcgov.github.io/PyRenew/
Apache License 2.0
14 stars 2 forks source link

Consider using `chex` in validation and/or testing #223

Open dylanhmorris opened 2 months ago

dylanhmorris commented 2 months ago

https://github.com/google-deepmind/chex

Could be helpful for improve, e.g. assertions about shapes:

https://github.com/CDCgov/multisignal-epi-inference/blob/4f203361a75be79c06f990b781194bece48d1f26/model/src/test/test_first_difference_ar.py#L29-L30

AFg6K7h4fhy2 commented 2 months ago

Beyond its formality and conciseness, it looks like it might be slightly faster, but I don't know if boost is trivial, or even if it's meaningful, given that this will be used in tests. Anyway, nice find.