bayinx 0.3.13__py3-none-any.whl → 0.3.14__py3-none-any.whl

bayinx/core/_optimization.py

@@ -1,4 +1,4 @@
-from typing import Any, Tuple
+from typing import Any, Tuple, TypeVar
 
 import equinox as eqx
 import jax.lax as lax
@@ -9,16 +9,16 @@ from optax import GradientTransformation, OptState, Schedule
 
 from ._model import Model
 
-
+M = TypeVar("M", bound=Model)
 @eqx.filter_jit
 def optimize_model(
-    model: Model,
+    model: M,
     max_iters: int,
     data: Any = None,
     learning_rate: float = 1,
     weight_decay: float = 0.0,
     tolerance: float = 1e-4,
-) -> Model:
+) -> M:
     """
     Optimize the dynamic parameters of the model.
 
@@ -35,9 +35,9 @@ def optimize_model(
     # Derive gradient for posterior
     @eqx.filter_jit
     @eqx.filter_grad
-    def eval_grad(dyn: Model):
+    def eval_grad(dyn: M):
         # Reconstruct model
-        model: Model = eqx.combine(dyn, static)
+        model: M = eqx.combine(dyn, static)
 
         # Evaluate posterior
         return model.eval(data)

bayinx/core/_parameter.py

@@ -5,8 +5,6 @@ import jax.tree as jt
 from jaxtyping import PyTree
 
 T = TypeVar("T", bound=PyTree)
-
-
 class Parameter(eqx.Module, Generic[T]):
     """
     A container for a parameter of a `Model`.

{bayinx-0.3.13.dist-info → bayinx-0.3.14.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bayinx
-Version: 0.3.13
+Version: 0.3.14
 Summary: Bayesian Inference with JAX
 License-File: LICENSE
 Requires-Python: >=3.12

{bayinx-0.3.13.dist-info → bayinx-0.3.14.dist-info}/RECORD

@@ -6,8 +6,8 @@ bayinx/core/__init__.py,sha256=Qmy0EjzqqKwI9F8rjmC9j6J8hiDw6A54yOck2WuQJkY,344
 bayinx/core/_constraint.py,sha256=Gx07ZT66VE2y-qZCmBDm3_y0wO4xQyslZW10Lec1_lM,761
 bayinx/core/_flow.py,sha256=3q4rKvATnbUpuj4ICUd4lIxu_3z7GRDuNujVhAku1X0,2342
 bayinx/core/_model.py,sha256=FJUyYVE9e2uTFamxtSMKY_VV2stiU2QF68Wl_7EAKEU,2895
-bayinx/core/_optimization.py,sha256=cXO07guDG5kd64kWq4_de-gLgxTT6vIOU3IOL3TMl6U,2583
-bayinx/core/_parameter.py,sha256=UfLonbTwxzr1g76Cf3HzAh9u4UBtlcaLByYLXgq-aCQ,1082
+bayinx/core/_optimization.py,sha256=dL3COMFTP0_FeD44hreNybh_UD6zqJy6sSP54ITJsBc,2605
+bayinx/core/_parameter.py,sha256=r20JedTW2lY0miNNh9y6LeIVAsGX1kP_rlGxphW_jZg,1080
 bayinx/core/_variational.py,sha256=bQYiN8c4AGPt4hsNT68zN7J6o0fdsLGwVjbDyl62LnI,5639
 bayinx/dists/__init__.py,sha256=BIrypqMnTLWK3a_zw8fYKMyuEMxP_qGsLfLeScias0o,118
 bayinx/dists/bernoulli.py,sha256=xMV9BgtVX_1XkPdZ43q0meMIEkgMyuUPx--dyo6_DKs,1006
@@ -31,7 +31,7 @@ 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.13.dist-info/METADATA,sha256=1mcHMTXrzMGPcibMy_vHdBJMNPtT0VUF83eVBS_MJlg,3080
-bayinx-0.3.13.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-bayinx-0.3.13.dist-info/licenses/LICENSE,sha256=VMhLhj5hx6VAENZBaNfXrmsNl7ov9uRh0jZ6D3ltgv4,1070
-bayinx-0.3.13.dist-info/RECORD,,
+bayinx-0.3.14.dist-info/METADATA,sha256=ba_0a1aYSrKIGhGFG_IWA7TltEmGpK78IqCFeLV7XVI,3080
+bayinx-0.3.14.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+bayinx-0.3.14.dist-info/licenses/LICENSE,sha256=VMhLhj5hx6VAENZBaNfXrmsNl7ov9uRh0jZ6D3ltgv4,1070
+bayinx-0.3.14.dist-info/RECORD,,