bayinx 0.2.24__py3-none-any.whl → 0.2.26__py3-none-any.whl

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/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/core/utils.py

@@ -1,54 +1 @@
-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

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.dist-info → bayinx-0.2.26.dist-info}/METADATA

@@ -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.dist-info → bayinx-0.2.26.dist-info}/RECORD

@@ -2,9 +2,9 @@ bayinx/__init__.py,sha256=l20JdkSsE_XGZlZFNEtySXf4NIlbjrao14vXPB-H6aQ,45
 bayinx/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 bayinx/core/__init__.py,sha256=7vW2F8t3K4TWlSu5nZrYCdUrz5N9FMIfQQBn3IoeH6o,150
 bayinx/core/constraints.py,sha256=Y8FJX3CkgnLQ3HXuTPGuzvLtXlKs0B7z0-YymoHgdfg,1682
-bayinx/core/flow.py,sha256=oZE0OHCninIHjp-WVLFWd1DaN0-qXxNWFAUAdgIDmRU,2423
-bayinx/core/model.py,sha256=t7s5Yt4E3iC_MPvynJnk6lb4OLal7Gnk59tIZ6e5M4I,2203
-bayinx/core/utils.py,sha256=-YewhqzMFL3GJEjVdm3LgaZyHwDs9IVYllU9wAXZrtw,1859
+bayinx/core/flow.py,sha256=9swS5wh7AsIZWgG_IhQS-upcPlr7G-juaP_5rsbX6G0,2363
+bayinx/core/model.py,sha256=U1xBnAXnIvFJjWF-XIT8BYjP9PtoRZY_PwyhRwJf-HA,2144
+bayinx/core/utils.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
 bayinx/core/variational.py,sha256=vUZ6u5CXCHfs6ZrA8PF4PHfmUXHTK2RJGHyZ3afFfsg,4820
 bayinx/dists/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 bayinx/dists/bernoulli.py,sha256=xMV9BgtVX_1XkPdZ43q0meMIEkgMyuUPx--dyo6_DKs,1006
@@ -19,10 +19,10 @@ bayinx/mhx/vi/meanfield.py,sha256=LNLwfjKO9os7YBmRBpGEiFxlxonuN7DHVFEmXV3hvfI,38
 bayinx/mhx/vi/normalizing_flow.py,sha256=nj7bpIoMJl6GTOXPxQCAsPArchbHd5vwwqMm3cLbMII,4791
 bayinx/mhx/vi/standard.py,sha256=HaJsIz70Qo1Ql2hMQ-GQhcnfWiOGtyxgkOsm_yQaDKI,1718
 bayinx/mhx/vi/flows/__init__.py,sha256=Hn0Wqvvyv8Vr-mFmimwgNKCByxj-fjrlIvdR7tUSolg,180
-bayinx/mhx/vi/flows/fullaffine.py,sha256=2QbOtA1Jmu-yRcJeFmCKc8N1atm8G7JXYMLEZaEXKV0,1711
+bayinx/mhx/vi/flows/fullaffine.py,sha256=TUcjXuDeFLHL9SYOLEU6kQSkEiyijztfBY2AAis7Pn0,2034
 bayinx/mhx/vi/flows/planar.py,sha256=u9ZVwEeOv4fMfwiORlseCz463atsWMuid_9crRg05Z8,1919
 bayinx/mhx/vi/flows/radial.py,sha256=c-SWybGn_jmgBQk9ZMQ5uHKPzcdhowNp8MD8t1-8VZM,2501
 bayinx/mhx/vi/flows/sylvester.py,sha256=ppK0BmS_ThvrCEhJiP_-p-kj67TQHSlU_RUZpDbIhsQ,469
-bayinx-0.2.24.dist-info/METADATA,sha256=sR0C0Pk5vrAmdvAtB3faXZO-hIDpKzqLjnXcfMsikjw,3058
-bayinx-0.2.24.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-bayinx-0.2.24.dist-info/RECORD,,
+bayinx-0.2.26.dist-info/METADATA,sha256=EReEQKQXgy71MtuqSgFRkn0YhVOUKT1hFECt12YEXxE,3058
+bayinx-0.2.26.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+bayinx-0.2.26.dist-info/RECORD,,