bayinx 0.2.10__tar.gz → 0.2.12__tar.gz

{bayinx-0.2.10 → bayinx-0.2.12}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bayinx
-Version: 0.2.10
+Version: 0.2.12
 Summary: Bayesian Inference with JAX
 Requires-Python: >=3.12
 Requires-Dist: equinox>=0.11.12

{bayinx-0.2.10 → bayinx-0.2.12}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "bayinx"
-version = "0.2.10"
+version = "0.2.12"
 description = "Bayesian Inference with JAX"
 readme = "README.md"
 requires-python = ">=3.12"
@@ -19,4 +19,4 @@ build-backend = "hatchling.build"
 addopts = "-q --benchmark-min-rounds=30 --benchmark-columns=rounds,mean,median,stddev --benchmark-group-by=func"
 
 [dependency-groups]
-dev = ["pytest>=8.3.5", "pytest-benchmark>=5.1.0"]
+dev = ["matplotlib>=3.10.1", "pytest>=8.3.5", "pytest-benchmark>=5.1.0"]

{bayinx-0.2.10 → bayinx-0.2.12}/src/bayinx/mhx/vi/flows/__init__.py

@@ -1,3 +1,3 @@
-from bayinx.mhx.vi.flows.affine import Affine as Affine
+from bayinx.mhx.vi.flows.fullaffine import FullAffine as FullAffine
 from bayinx.mhx.vi.flows.planar import Planar as Planar
 from bayinx.mhx.vi.flows.radial import Radial as Radial

bayinx-0.2.10/src/bayinx/mhx/vi/flows/affine.py → bayinx-0.2.12/src/bayinx/mhx/vi/flows/fullaffine.py

@@ -9,7 +9,7 @@ from jaxtyping import Array, Float, Scalar
 from bayinx.core import Flow
 
 
-class Affine(Flow):
+class FullAffine(Flow):
     """
     An affine flow.
 

{bayinx-0.2.10 → bayinx-0.2.12}/src/bayinx/mhx/vi/flows/planar.py

@@ -30,16 +30,16 @@ class Planar(Flow):
         - `dim`: The dimension of the parameter space.
         """
         self.params = {
-            "u": jr.normal(key, (dim,)),
-            "w": jr.normal(key, (dim,)),
-            "b": jr.normal(key, (1,)),
+            "u": jnp.ones(dim),
+            "w": jnp.ones(dim),
+            "b": jnp.zeros(1),
         }
         self.constraints = {}
 
     @eqx.filter_jit
     @partial(jax.vmap, in_axes=(None, 0))
     def forward(self, draws: Array) -> Array:
-        params = self.constrain_pars()
+        params = self.transform_pars()
 
         # Extract parameters
         w: Array = params["w"]
@@ -54,7 +54,7 @@ class Planar(Flow):
     @eqx.filter_jit
     @partial(jax.vmap, in_axes=(None, 0))
     def adjust_density(self, draws: Array) -> Tuple[Scalar, Array]:
-        params = self.constrain_pars()
+        params = self.transform_pars()
 
         # Extract parameters
         w: Array = params["w"]

{bayinx-0.2.10 → bayinx-0.2.12}/src/bayinx/mhx/vi/normalizing_flow.py

@@ -115,7 +115,7 @@ class NormalizingFlow(Variational):
         return filter_spec
 
     @eqx.filter_jit
-    def elbo(self, n: int, key: Key, data: Any = None) -> Scalar:
+    def elbo(self, n: int, key: Key = jr.PRNGKey(0), data: Any = None) -> Scalar:
         dyn, static = eqx.partition(self, self.filter_spec())
 
         @eqx.filter_jit

{bayinx-0.2.10 → bayinx-0.2.12}/tests/test_variational.py

@@ -8,7 +8,7 @@ from jaxtyping import Array
 from bayinx import Model
 from bayinx.dists import normal
 from bayinx.mhx.vi import MeanField, NormalizingFlow, Standard
-from bayinx.mhx.vi.flows import Affine, Planar, Radial
+from bayinx.mhx.vi.flows import FullAffine, Planar, Radial
 
 
 # Tests ----
@@ -44,7 +44,7 @@ def test_meanfield(benchmark, var_draws):
         vari.fit(10000, var_draws=var_draws)
 
     benchmark(benchmark_fit)
-    vari = vari.fit(10000)
+    vari = vari.fit(20000)
 
     # Assert parameters are roughly correct
     assert all(abs(10.0 - vari.var_params["mean"]) < 0.1) and all(
@@ -77,14 +77,14 @@ def test_affine(benchmark, var_draws):
     model = NormalDist()
 
     # Construct normalizing flow variational
-    vari = NormalizingFlow(Standard(model), [Affine(2)], model)
+    vari = NormalizingFlow(Standard(model), [FullAffine(2)], model)
 
     # Optimize variational distribution
     def benchmark_fit():
         vari.fit(10000, var_draws=var_draws)
 
     benchmark(benchmark_fit)
-    vari = vari.fit(10000)
+    vari = vari.fit(20000)
 
     params = vari.flows[0].constrain_pars()
     assert (abs(10.0 - vari.flows[0].params["shift"]) < 0.1).all() and (
@@ -118,7 +118,7 @@ def test_flows(benchmark, var_draws):
 
     # Construct normalizing flow variational
     vari = NormalizingFlow(
-        Standard(model), [Planar(2), Radial(2), Planar(2), Radial(2), Planar(2)], model
+        Standard(model), [FullAffine(2), Planar(2), Radial(2)], model
     )
 
     # Optimize variational distribution
@@ -126,7 +126,7 @@ def test_flows(benchmark, var_draws):
         vari.fit(10000, var_draws=var_draws)
 
     benchmark(benchmark_fit)
-    vari = vari.fit(10000)
+    vari = vari.fit(20000)
 
     mean = vari.sample(1000).mean(0)
     var = vari.sample(1000).var(0)