bayinx 0.3.0__py3-none-any.whl → 0.3.1__py3-none-any.whl

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

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

@@ -5,8 +5,8 @@ bayinx/constraints/lower.py,sha256=wkYnWjaAEGQeXKfBo_gY0pcK9ElJUMkzGdAmWI8ykCk,1
 bayinx/core/__init__.py,sha256=jSwEFdXqi-Bj_X8_H-YuaXp5ebEQpZTG2T18zpquzPo,207
 bayinx/core/constraint.py,sha256=F6-TXQjzt-tcNm8bHkRcGEtyE9bZQf2RbAh_MKDuM20,760
 bayinx/core/flow.py,sha256=lAPJdQnrIxC3JoowTp77Gvm0p0v_xQn8FMjFJWMnWbc,2340
-bayinx/core/model.py,sha256=QnJUKaR6d5RCe_WIxD2oJtI8NJyFKWUWyCRVwOm0j3s,2276
-bayinx/core/parameter.py,sha256=fdyzun6TDnXxQT_KlarIJvWzn9n8bQgzfiVjWIIHk6k,998
+bayinx/core/model.py,sha256=RKAtsLXc6xDnXWz5upmx6Vz6JOoorw4WTfxTA7B7Lmg,2294
+bayinx/core/parameter.py,sha256=oxCCZcZ-DDBvfWzexfhQkSJPxNQnE1vYXtBhiEZG2bM,1025
 bayinx/core/variational.py,sha256=lqENISRrKY8ODLtl0D-D7TAA2gD7HGh37BnROM7p5hI,4783
 bayinx/dists/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 bayinx/dists/bernoulli.py,sha256=xMV9BgtVX_1XkPdZ43q0meMIEkgMyuUPx--dyo6_DKs,1006
@@ -25,6 +25,6 @@ bayinx/mhx/vi/flows/fullaffine.py,sha256=11y_A0oO3bkKDSz-EQ6Sf4Ec2M7vHZxw94EdvAD
 bayinx/mhx/vi/flows/planar.py,sha256=2I2WzIskl8MRpJkK13FQE3vSF-077qo8gRed_EL1Pn8,1920
 bayinx/mhx/vi/flows/radial.py,sha256=e0GfuO-CL8SVr3YnEfsxStpyKcJlFpzMyjMp3sa38hg,2503
 bayinx/mhx/vi/flows/sylvester.py,sha256=ppK0BmS_ThvrCEhJiP_-p-kj67TQHSlU_RUZpDbIhsQ,469
-bayinx-0.3.0.dist-info/METADATA,sha256=RLbnLgyMmnEh2BJmqex3MMWFFS3HgSU9NEeQEvkyfC0,3057
-bayinx-0.3.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-bayinx-0.3.0.dist-info/RECORD,,
+bayinx-0.3.1.dist-info/METADATA,sha256=Slp2nxR8HISCwCqIXY2El3GgqsO1v9_UbVeJq726w7k,3057
+bayinx-0.3.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+bayinx-0.3.1.dist-info/RECORD,,