bayinx 0.3.13__tar.gz → 0.3.14__tar.gz

{bayinx-0.3.13 → bayinx-0.3.14}/PKG-INFO

@@ -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 → bayinx-0.3.14}/pyproject.toml

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

{bayinx-0.3.13 → bayinx-0.3.14}/src/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-0.3.13 → bayinx-0.3.14}/src/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`.