bayinx 0.2.9__py3-none-any.whl → 0.2.11__py3-none-any.whl

bayinx/core/variational.py

@@ -89,6 +89,7 @@ class Variational(eqx.Module):
             max_iters: int,
             data: Any = None,
             learning_rate: float = 1,
+            weight_decay: float = 1e-4,
             tolerance: float = 1e-4,
             var_draws: int = 1,
             key: Key = jr.PRNGKey(0),
@@ -114,7 +115,7 @@ class Variational(eqx.Module):
 
             # Initialize optimizer
             optim: GradientTransformation = opx.chain(
-                opx.scale(-1.0), opx.nadam(schedule)
+                opx.scale(-1.0), opx.nadamw(schedule,weight_decay=weight_decay)
             )
             opt_state: OptState = optim.init(dyn)
 

bayinx/mhx/vi/flows/__init__.py

@@ -1,3 +1,3 @@
-from bayinx.mhx.vi.flows.affine import Affine as Affine
+from bayinx.mhx.vi.flows.fullaffine import FullAffine as FullAffine
 from bayinx.mhx.vi.flows.planar import Planar as Planar
 from bayinx.mhx.vi.flows.radial import Radial as Radial

bayinx/mhx/vi/flows/{affine.py → fullaffine.py}

@@ -9,7 +9,7 @@ from jaxtyping import Array, Float, Scalar
 from bayinx.core import Flow
 
 
-class Affine(Flow):
+class FullAffine(Flow):
     """
     An affine flow.
 

bayinx/mhx/vi/flows/planar.py

@@ -3,6 +3,7 @@ from typing import Callable, Dict, Tuple
 
 import equinox as eqx
 import jax
+import jax.nn as jnn
 import jax.numpy as jnp
 import jax.random as jr
 from jaxtyping import Array, Float, Scalar
@@ -36,10 +37,23 @@ class Planar(Flow):
         }
         self.constraints = {}
 
+    def transform_pars(self):
+        params = self.params
+
+        u = params['u']
+        w = params['w']
+        b = params['b']
+
+        m = jnn.softplus(w.dot(u)) - 1.0
+
+        u = u + (m - w.dot(u)) * w / (w**2).sum()
+
+        return {'u': u, 'w': w, 'b': b}
+
     @eqx.filter_jit
     @partial(jax.vmap, in_axes=(None, 0))
     def forward(self, draws: Array) -> Array:
-        params = self.constrain_pars()
+        params = self.transform_pars()
 
         # Extract parameters
         w: Array = params["w"]
@@ -54,7 +68,7 @@ class Planar(Flow):
     @eqx.filter_jit
     @partial(jax.vmap, in_axes=(None, 0))
     def adjust_density(self, draws: Array) -> Tuple[Scalar, Array]:
-        params = self.constrain_pars()
+        params = self.transform_pars()
 
         # Extract parameters
         w: Array = params["w"]

bayinx/mhx/vi/normalizing_flow.py

@@ -115,7 +115,7 @@ class NormalizingFlow(Variational):
         return filter_spec
 
     @eqx.filter_jit
-    def elbo(self, n: int, key: Key, data: Any = None) -> Scalar:
+    def elbo(self, n: int, key: Key = jr.PRNGKey(0), data: Any = None) -> Scalar:
         dyn, static = eqx.partition(self, self.filter_spec())
 
         @eqx.filter_jit

{bayinx-0.2.9.dist-info → bayinx-0.2.11.dist-info}/METADATA

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

{bayinx-0.2.9.dist-info → bayinx-0.2.11.dist-info}/RECORD

@@ -4,7 +4,7 @@ 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=yhraTVlNOJaU1NEYVrWpUXVzzWvY1Mq9ZOZv6V0_Vo0,5183
+bayinx/core/variational.py,sha256=3CsDyQkq1XgV2ZBLzGrm5XgUFoJBnT6glHDgxHNcbTc,5250
 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
@@ -14,13 +14,13 @@ bayinx/dists/normal.py,sha256=e9gXXAHeZQKjBndW2TnMvP3gtmvpfYGG7kehcpGeAoU,2590
 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=V0-R2Nc_3Vy5c2rkjDOwIutA0G-_kN6trxjwsT5FgsA,4774
+bayinx/mhx/vi/normalizing_flow.py,sha256=XBUWYZpm_Ipi6X9oTnGhqIs3ARY-5QFiuxM7uAWFRps,4790
 bayinx/mhx/vi/standard.py,sha256=m5gtcHfrYzV28h-Red3Zn6SxEgJlndeIXiIG5gDPecU,1703
-bayinx/mhx/vi/flows/__init__.py,sha256=V_Ng5cecKlLlFSI9ncmaiyvoy_d2EAfeDhBFcy5aQhA,168
-bayinx/mhx/vi/flows/affine.py,sha256=a205nNx6KRvOwGlnjI6YeDo7OTWPPIxffGZfAcTecNA,1707
-bayinx/mhx/vi/flows/planar.py,sha256=0BGdMm-GpTCJnxq9cOrgLl8IsHgGIL0eSFagWJNVdqQ,1944
+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=u4heNqxpmfXsACsE3RH8XkBd04Emd8G674M3sSd7CxM,2232
 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.9.dist-info/METADATA,sha256=xp6L_DdXPC-TMHV4SL5LdIuhzX8GizUlx2muMgSFcy0,3057
-bayinx-0.2.9.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-bayinx-0.2.9.dist-info/RECORD,,
+bayinx-0.2.11.dist-info/METADATA,sha256=TkIGXb5baSFeP5K4s1bynQ3QqNsR-DKePlAOV_vV0Nc,3058
+bayinx-0.2.11.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+bayinx-0.2.11.dist-info/RECORD,,