pyro-ppl / numpyro / 13773109953

push

github

Flax NNX integration (#1990)

* This is a combination of 9 commits.

init

fix

fix

fix lint

rm files

rm files

rm files

rm files

rm files

* improvement

* improvements

* add more tests

* add mcmc cases

* Refactor

add example

rm test patch

feedback 1

feedback 2

feedback 3

improvements

fix test

hacky way to find batch normalization layers

remove unused code

simplyfy code

modularize

tests

simplify

simplify

* support python 3.9

* init example

* fix model

* cleanup nb

* artifacts docs

* remove code

eager inizialization approach

rm unused code 1

rm unused code 2

rm code

rm code

rm  code

rm code

* rm code

* refactor

* skip tests

* redback 1/n

* redback 2/n

* lint

* skip test

* feedback 3/n

* partial

* feedback states

* clean tests part 1

* set priors

* better split

* rm mutable from signature

* fix numpyro_mutable

* fix test

* feedback part1

* try fix apply function

* Revert "try fix apply function"

This reverts commit 17db8c14c.

* Revert "Revert "try fix apply function""

This reverts commit 9a870d66d.

* rm reduntand line

* better test name

* are we updating more than needed?

* rename and bring back update after model call

* add prior component to example

* finalize docs

* to_dict

* stop gradient

* finall comments

41 of 43 new or added lines in 1 file covered. (95.35%)

13468 of 14506 relevant lines covered (92.84%)

1.99 hits per line