bayinx 0.3.0__tar.gz → 0.3.1__tar.gz

{bayinx-0.3.0 → bayinx-0.3.1}/PKG-INFO

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

{bayinx-0.3.0 → bayinx-0.3.1}/pyproject.toml

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

{bayinx-0.3.0 → bayinx-0.3.1}/src/bayinx/core/model.py

@@ -1,25 +1,25 @@
 from abc import abstractmethod
-from typing import Any, Dict, Tuple
+from typing import Any, Dict, Generic, Tuple, TypeVar
 
 import equinox as eqx
 import jax.numpy as jnp
 import jax.tree as jt
-from jaxtyping import Array, PyTree, Scalar
+from jaxtyping import PyTree, Scalar
 
 from bayinx.core.constraint import Constraint
 from bayinx.core.parameter import Parameter
 
-
-class Model(eqx.Module):
+T = TypeVar('T', bound=PyTree)
+class Model(eqx.Module, Generic[T]):
     """
     An abstract base class used to define probabilistic models.
 
     # Attributes
-    - `params`: A dictionary of Arrays or PyTrees containing Arrays representing parameters of the model.
+    - `params`: A dictionary of parameters.
     - `constraints`: A dictionary of constraints.
     """
 
-    params: Dict[str, Parameter]
+    params: Dict[str, Parameter[T]]
     constraints: Dict[str, Constraint]
 
     @abstractmethod
@@ -47,14 +47,14 @@ class Model(eqx.Module):
 
     # Add constrain method
     @eqx.filter_jit
-    def constrain_params(self) -> Tuple[Dict[str, Parameter], Scalar]:
+    def constrain_params(self) -> Tuple[Dict[str, Parameter[T]], Scalar]:
         """
         Constrain `params` to the appropriate domain.
 
         # Returns
         A dictionary of PyTrees representing the constrained parameters and the adjustment to the posterior density.
         """
-        t_params: Dict[str, Array | PyTree] = self.params
+        t_params: Dict[str, Parameter[T]] = self.params
         target: Scalar = jnp.array(0.0)
 
         for par, map in self.constraints.items():
@@ -67,7 +67,8 @@ class Model(eqx.Module):
         return t_params, target
 
     # Add default transform method
-    def transform_params(self) -> Tuple[Dict[str, Parameter], Scalar]:
+    @eqx.filter_jit
+    def transform_params(self) -> Tuple[Dict[str, Parameter[T]], Scalar]:
         """
         Apply a custom transformation to `params` if needed.
 

{bayinx-0.3.0 → bayinx-0.3.1}/src/bayinx/core/parameter.py

@@ -1,11 +1,11 @@
-from typing import Self
+from typing import Generic, Self, TypeVar
 
 import equinox as eqx
 import jax.tree as jt
-from jaxtyping import Array, PyTree
+from jaxtyping import PyTree
 
-
-class Parameter(eqx.Module):
+T = TypeVar('T', bound=PyTree)
+class Parameter(eqx.Module, Generic[T]):
     """
     A container for a parameter of a `Model`.
 
@@ -14,10 +14,10 @@ class Parameter(eqx.Module):
     # Attributes
     - `vals`: The parameter's value(s).
     """
-    vals: Array | PyTree
+    vals: T
 
 
-    def __init__(self, values: Array | PyTree):
+    def __init__(self, values: T):
         # Insert parameter values
         self.vals = values
 

{bayinx-0.3.0 → bayinx-0.3.1}/tests/test_variational.py

@@ -1,7 +1,10 @@
 
+from typing import Dict
+
 import equinox as eqx
 import jax.numpy as jnp
 import pytest
+from jaxtyping import Array
 
 from bayinx import Model, Parameter
 from bayinx.dists import normal
@@ -13,7 +16,9 @@ from bayinx.mhx.vi.flows import FullAffine, Planar, Radial
 @pytest.mark.parametrize("var_draws", [1, 10, 100])
 def test_meanfield(benchmark, var_draws):
     # Construct model definition
-    class NormalDist(Model):
+    class NormalDist(Model[Array]):
+        params: Dict[str, Parameter[Array]]
+
         def __init__(self):
             self.params = {"mu": Parameter(jnp.array([0.0, 0.0]))}
             self.constraints = {}