advi_jax
tests
