bayinx 0.2.24__tar.gz → 0.2.26__tar.gz

{bayinx-0.2.24 → bayinx-0.2.26}/PKG-INFO

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

{bayinx-0.2.24 → bayinx-0.2.26}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "bayinx"
-version = "0.2.24"
+version = "0.2.26"
 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.24"
+current_version = "0.2.26"
 parse = "(?P<major>\\d+)\\.(?P<minor>\\d+)\\.(?P<patch>\\d+)"
 serialize = ["{major}.{minor}.{patch}"]
 search = "{current_version}"

{bayinx-0.2.24 → bayinx-0.2.26}/src/bayinx/core/flow.py

@@ -5,10 +5,8 @@ import equinox as eqx
 import jax.tree_util as jtu
 from jaxtyping import Array, Float
 
-from bayinx.core.utils import __MyMeta
 
-
-class Flow(eqx.Module, metaclass=__MyMeta):
+class Flow(eqx.Module):
     """
     A superclass used to define continuously parameterized diffeomorphisms for normalizing flows.
 

{bayinx-0.2.24 → bayinx-0.2.26}/src/bayinx/core/model.py

@@ -7,10 +7,9 @@ import jax.tree_util as jtu
 from jaxtyping import Array, Scalar
 
 from bayinx.core.constraints import Constraint
-from bayinx.core.utils import __MyMeta
 
 
-class Model(eqx.Module, metaclass=__MyMeta):
+class Model(eqx.Module):
     """
     A superclass used to define probabilistic models.
 

bayinx-0.2.26/src/bayinx/mhx/__init__.py

@@ -0,0 +1 @@
+

{bayinx-0.2.24 → bayinx-0.2.26}/src/bayinx/mhx/vi/flows/fullaffine.py

@@ -33,7 +33,22 @@ 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)
+
+
+        return params
+
+
 
     @eqx.filter_jit
     def forward(self, draws: Array) -> Array:

bayinx-0.2.24/src/bayinx/core/utils.py

@@ -1,54 +0,0 @@
-from typing import Callable, Dict
-
-import equinox as eqx
-from jaxtyping import Array
-
-
-class __MyMeta(type(eqx.Module)):
-    """
-    Metaclass to ensure attribute types are respected.
-    """
-
-    def __call__(cls, *args, **kwargs):
-        obj = super().__call__(*args, **kwargs)
-
-        # Check parameters are a Dict of JAX Arrays
-        if not isinstance(obj.params, Dict):
-            raise ValueError(
-                f"Model {cls.__name__} must initialize 'params' as a dictionary."
-            )
-
-        for key, value in obj.params.items():
-            if not isinstance(value, Array):
-                raise TypeError(f"Parameter '{key}' must be a JAX Array.")
-
-        # Check constraints are a Dict of functions
-        if not isinstance(obj.constraints, Dict):
-            raise ValueError(
-                f"Model {cls.__name__} must initialize 'constraints' as a dictionary."
-            )
-
-        for key, value in obj.constraints.items():
-            if not isinstance(value, Callable):
-                raise TypeError(f"Constraint for parameter '{key}' must be a function.")
-
-        # Check that the constrain method returns a dict equivalent to `params`
-        t_params: Dict[str, Array] = obj.constrain_pars()
-
-        if not isinstance(t_params, Dict):
-            raise ValueError(
-                f"The 'constrain' method of {cls.__name__} must return a Dict of JAX Arrays."
-            )
-
-        for key, value in t_params.items():
-            if not isinstance(value, Array):
-                raise TypeError(f"Constrained parameter '{key}' must be a JAX Array.")
-
-            if not value.shape == obj.params[key].shape:
-                raise ValueError(
-                    f"Constrained parameter '{key}' must have same shape as unconstrained counterpart."
-                )
-
-        # Check transform_pars
-
-        return obj