pyro-ppl / numpyro / 14120810158

92%

push

github

Equinox Integration (#2005)

* added initial support for equinox models

* removed accidental line spaces

* used alias for jax.tree_util in contrib/module

* added docstrings for update_tree_params

* working eqx_module that manages state internally

* fixed issue with state getting broadcasted

* added equinox as dev dependency

* made eqx_module functional

* better docstrings for equinox modules

* changed eqx smoke test to use autodelta

* stopped gradient for eqx state

* updated docstring to reflect stopping eqx state gradient

* removed update_tree_params and switched random_eqx_module to use _update_params

* added docs for eqx_modules, mutable. stopped gradient on mutable_state in svi loss_fn

* fixed docs

* adjusted test to not have latent x in model

* removed stopping state gradient in svi loss_fn

* updated docstring to reflect no gradient stopping

---------

Co-authored-by: kylejcaron <kylejcaron@bitbucket.org>

29 of 33 new or added lines in 1 file covered. (87.88%)

13619 of 14665 relevant lines covered (92.87%)

1.98 hits per line