bayinx 0.2.25__tar.gz → 0.2.27__tar.gz

{bayinx-0.2.25 → bayinx-0.2.27}/PKG-INFO

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

{bayinx-0.2.25 → bayinx-0.2.27}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "bayinx"
-version = "0.2.25"
+version = "0.2.27"
 description = "Bayesian Inference with JAX"
 readme = "README.md"
 requires-python = ">=3.12"
@@ -19,7 +19,7 @@ build-backend = "hatchling.build"
 addopts = "-q --benchmark-min-rounds=30 --benchmark-columns=rounds,mean,median,stddev --benchmark-group-by=func"
 
 [tool.bumpversion]
-current_version = "0.2.25"
+current_version = "0.2.27"
 parse = "(?P<major>\\d+)\\.(?P<minor>\\d+)\\.(?P<patch>\\d+)"
 serialize = ["{major}.{minor}.{patch}"]
 search = "{current_version}"

{bayinx-0.2.25 → bayinx-0.2.27}/src/bayinx/mhx/vi/flows/fullaffine.py

@@ -33,11 +33,26 @@ class FullAffine(Flow):
             "scale": jnp.zeros((dim, dim)),
         }
 
-        self.constraints = {"scale": lambda m: jnp.tril(jnp.exp(m))}
+        self.constraints = {"scale": lambda m: jnp.tril(m)}
+
+    def transform_pars(self):
+        # Get constrained parameters
+        params = self.constrain_pars()
+
+        # Extract diagonal and apply exponential
+        diag: Array = jnp.exp(jnp.diag(params['scale']))
+
+        # Fill diagonal
+        params['scale'] = jnp.fill_diagonal(params['scale'], diag, inplace=False)
+
+
+        return params
+
+
 
     @eqx.filter_jit
     def forward(self, draws: Array) -> Array:
-        params = self.constrain_pars()
+        params = self.transform_pars()
 
         # Extract parameters
         shift: Array = params["shift"]
@@ -51,7 +66,7 @@ class FullAffine(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
         shift: Array = params["shift"]

{bayinx-0.2.25 → bayinx-0.2.27}/tests/test_variational.py

@@ -88,7 +88,7 @@ def test_affine(benchmark, var_draws):
     benchmark(benchmark_fit)
     vari = vari.fit(20000,var_draws=var_draws)
 
-    params = vari.flows[0].constrain_pars()
+    params = vari.flows[0].transform_pars()
     assert (abs(10.0 - vari.flows[0].params["shift"]) < 0.1).all() and (
         abs(jnp.eye(2) - params["scale"]) < 0.1
     ).all()