bayinx 0.2.12__tar.gz → 0.2.18__tar.gz

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

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

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

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

{bayinx-0.2.12 → bayinx-0.2.18}/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.18}/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.18}/src/bayinx/mhx/vi/standard.py

@@ -52,7 +52,7 @@ class Standard(Variational):
             x=draws,
             mu=jnp.array(0.0),
             sigma=jnp.array(1.0),
-        ).sum(axis=1)
+        ).sum(axis=1, keepdims=True)
 
     @eqx.filter_jit
     def filter_spec(self):