bayinx 0.2.12__py3-none-any.whl → 0.2.18__py3-none-any.whl

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/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/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):

{bayinx-0.2.12.dist-info → bayinx-0.2.18.dist-info}/METADATA

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

@@ -4,23 +4,23 @@ bayinx/core/__init__.py,sha256=7vW2F8t3K4TWlSu5nZrYCdUrz5N9FMIfQQBn3IoeH6o,150
 bayinx/core/flow.py,sha256=oZE0OHCninIHjp-WVLFWd1DaN0-qXxNWFAUAdgIDmRU,2423
 bayinx/core/model.py,sha256=-rT3NHjxqGB0lDBMi0Mr9XNOz1_TUnJWtd4ITj0rsus,2257
 bayinx/core/utils.py,sha256=-YewhqzMFL3GJEjVdm3LgaZyHwDs9IVYllU9wAXZrtw,1859
-bayinx/core/variational.py,sha256=3CsDyQkq1XgV2ZBLzGrm5XgUFoJBnT6glHDgxHNcbTc,5250
+bayinx/core/variational.py,sha256=k9wWn7Tnj3eET-qK1pZtzDyPZVvQTRUexJUBVSdGXOA,5251
 bayinx/dists/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 bayinx/dists/bernoulli.py,sha256=xMV9BgtVX_1XkPdZ43q0meMIEkgMyuUPx--dyo6_DKs,1006
 bayinx/dists/binomial.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 bayinx/dists/gamma.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 bayinx/dists/gamma2.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-bayinx/dists/normal.py,sha256=e9gXXAHeZQKjBndW2TnMvP3gtmvpfYGG7kehcpGeAoU,2590
+bayinx/dists/normal.py,sha256=OOKg46y5hHFP76ydbRjEXaDkgefZcj9sd0XAl7yokww,2587
 bayinx/mhx/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
 bayinx/mhx/vi/__init__.py,sha256=YfkXKsqo9Dk_AmQGjZKm4vfG8eLer2ez92G-cOExphs,193
 bayinx/mhx/vi/meanfield.py,sha256=LNLwfjKO9os7YBmRBpGEiFxlxonuN7DHVFEmXV3hvfI,3876
 bayinx/mhx/vi/normalizing_flow.py,sha256=XBUWYZpm_Ipi6X9oTnGhqIs3ARY-5QFiuxM7uAWFRps,4790
-bayinx/mhx/vi/standard.py,sha256=m5gtcHfrYzV28h-Red3Zn6SxEgJlndeIXiIG5gDPecU,1703
+bayinx/mhx/vi/standard.py,sha256=HaJsIz70Qo1Ql2hMQ-GQhcnfWiOGtyxgkOsm_yQaDKI,1718
 bayinx/mhx/vi/flows/__init__.py,sha256=Hn0Wqvvyv8Vr-mFmimwgNKCByxj-fjrlIvdR7tUSolg,180
 bayinx/mhx/vi/flows/fullaffine.py,sha256=2QbOtA1Jmu-yRcJeFmCKc8N1atm8G7JXYMLEZaEXKV0,1711
 bayinx/mhx/vi/flows/planar.py,sha256=qmtWpIBXRct2seI78pkmtF0X7cASUBELqmZmf2QS5Gs,1918
 bayinx/mhx/vi/flows/radial.py,sha256=c-SWybGn_jmgBQk9ZMQ5uHKPzcdhowNp8MD8t1-8VZM,2501
 bayinx/mhx/vi/flows/sylvester.py,sha256=ppK0BmS_ThvrCEhJiP_-p-kj67TQHSlU_RUZpDbIhsQ,469
-bayinx-0.2.12.dist-info/METADATA,sha256=q4e6XXwZ6ejyBWsyk_wXGDqJG9YCBK1gew93Pg_PncU,3058
-bayinx-0.2.12.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-bayinx-0.2.12.dist-info/RECORD,,
+bayinx-0.2.18.dist-info/METADATA,sha256=849aiZhV6f578GKb62AfN_YgIoyJZ7FfA-fYw-Ogk_0,3058
+bayinx-0.2.18.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+bayinx-0.2.18.dist-info/RECORD,,