bayinx 0.2.12__tar.gz → 0.2.13__tar.gz

{bayinx-0.2.12 → bayinx-0.2.13}/PKG-INFO

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

{bayinx-0.2.12 → bayinx-0.2.13}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "bayinx"
-version = "0.2.12"
+version = "0.2.13"
 description = "Bayesian Inference with JAX"
 readme = "README.md"
 requires-python = ">=3.12"

{bayinx-0.2.12 → bayinx-0.2.13}/src/bayinx/core/variational.py

@@ -115,7 +115,7 @@ class Variational(eqx.Module):
 
             # Initialize optimizer
             optim: GradientTransformation = opx.chain(
-                opx.scale(-1.0), opx.nadamw(schedule,weight_decay=weight_decay)
+                opx.scale(-1.0), opx.nadamw(schedule, weight_decay=weight_decay)
             )
             opt_state: OptState = optim.init(dyn)
 

{bayinx-0.2.12 → bayinx-0.2.13}/src/bayinx/dists/normal.py

@@ -1,17 +1,12 @@
-# MARK: Imports ----
 import jax.lax as _lax
+from jaxtyping import Array, ArrayLike, Float, Real
 
-## Typing
-from jaxtyping import Array, Real
-
-# MARK: Constants
 _PI = 3.141592653589793
 
 
-# MARK: Functions ----
 def prob(
-    x: Real[Array, "..."], mu: Real[Array, "..."], sigma: Real[Array, "..."]
-) -> Real[Array, "..."]:
+    x: Real[ArrayLike, "..."], mu: Real[ArrayLike, "..."], sigma: Real[ArrayLike, "..."]
+) -> Float[Array, "..."]:
     """
     The probability density function (PDF) for a Normal distribution.
 
@@ -30,8 +25,8 @@ def prob(
 
 
 def logprob(
-    x: Real[Array, "..."], mu: Real[Array, "..."], sigma: Real[Array, "..."]
-) -> Real[Array, "..."]:
+    x: Real[ArrayLike, "..."], mu: Real[ArrayLike, "..."], sigma: Real[ArrayLike, "..."]
+) -> Float[Array, "..."]:
     """
     The log of the probability density function (log PDF) for a Normal distribution.
 
@@ -48,8 +43,8 @@ def logprob(
 
 
 def uprob(
-    x: Real[Array, "..."], mu: Real[Array, "..."], sigma: Real[Array, "..."]
-) -> Real[Array, "..."]:
+    x: Real[ArrayLike, "..."], mu: Real[ArrayLike, "..."], sigma: Real[ArrayLike, "..."]
+) -> Float[Array, "..."]:
     """
     The unnormalized probability density function (uPDF) for a Normal distribution.
 
@@ -66,8 +61,8 @@ def uprob(
 
 
 def ulogprob(
-    x: Real[Array, "..."], mu: Real[Array, "..."], sigma: Real[Array, "..."]
-) -> Real[Array, "..."]:
+    x: Real[ArrayLike, "..."], mu: Real[ArrayLike, "..."], sigma: Real[ArrayLike, "..."]
+) -> Float[Array, "..."]:
     """
     The log of the unnormalized probability density function (log uPDF) for a Normal distribution.
 

{bayinx-0.2.12 → bayinx-0.2.13}/src/bayinx/mhx/vi/normalizing_flow.py

@@ -1,6 +1,8 @@
+from functools import partial
 from typing import Any, Callable, Self, Tuple
 
 import equinox as eqx
+import jax
 import jax.flatten_util as jfu
 import jax.numpy as jnp
 import jax.random as jr
@@ -59,6 +61,7 @@ class NormalizingFlow(Variational):
         return draws
 
     @eqx.filter_jit
+    @partial(jax.vmap, in_axes=(None, 0))
     def eval(self, draws: Array) -> Array:
         # Evaluate base density
         variational_evals: Array = self.base.eval(draws)