bayinx 0.2.10__tar.gz → 0.2.11__tar.gz

{bayinx-0.2.10 → bayinx-0.2.11}/PKG-INFO

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

{bayinx-0.2.10 → bayinx-0.2.11}/pyproject.toml

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

{bayinx-0.2.10 → bayinx-0.2.11}/src/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-0.2.10/src/bayinx/mhx/vi/flows/affine.py → bayinx-0.2.11/src/bayinx/mhx/vi/flows/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-0.2.10 → bayinx-0.2.11}/src/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-0.2.10 → bayinx-0.2.11}/src/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.10 → bayinx-0.2.11}/tests/test_variational.py

@@ -8,7 +8,7 @@ from jaxtyping import Array
 from bayinx import Model
 from bayinx.dists import normal
 from bayinx.mhx.vi import MeanField, NormalizingFlow, Standard
-from bayinx.mhx.vi.flows import Affine, Planar, Radial
+from bayinx.mhx.vi.flows import FullAffine, Planar
 
 
 # Tests ----
@@ -77,7 +77,7 @@ def test_affine(benchmark, var_draws):
     model = NormalDist()
 
     # Construct normalizing flow variational
-    vari = NormalizingFlow(Standard(model), [Affine(2)], model)
+    vari = NormalizingFlow(Standard(model), [FullAffine(2)], model)
 
     # Optimize variational distribution
     def benchmark_fit():
@@ -95,30 +95,37 @@ def test_affine(benchmark, var_draws):
 @pytest.mark.parametrize("var_draws", [1, 10, 100])
 def test_flows(benchmark, var_draws):
     # Construct model definition
-    class NormalDist(Model):
+    class Banana(Model):
         params: Dict[str, Array]
         constraints: Dict[str, Callable[[Array], Array]]
 
         def __init__(self):
-            self.params = {"mu": jnp.array([0.0, 0.0])}
+            self.params = {
+                'x': jnp.array(0.0),
+                'y': jnp.array(0.0)
+            }
             self.constraints = {}
 
-        @eqx.filter_jit
-        def eval(self, data: dict):
-            # Get constrained parameters
-            params = self.constrain_pars()
+        def eval(self, data = None):
+            params: Dict[str, Array] = self.params
+            # Extract parameters
+            x: Array = params['x']
+            y: Array = params['y']
 
-            # Evaluate mu ~ N(10,1)
-            return jnp.sum(
-                normal.logprob(x=params["mu"], mu=jnp.array(10.0), sigma=jnp.array(1.0))
-            )
+            # Initialize target density
+            target = jnp.array(0.0)
+
+            target += normal.logprob(x, mu = jnp.array(0.0), sigma = jnp.array(1.0))
+            target += normal.logprob(y, mu = x**2 + x, sigma = jnp.array(1.0))
+
+            return target
 
     # Construct model
-    model = NormalDist()
+    model = Banana()
 
     # Construct normalizing flow variational
     vari = NormalizingFlow(
-        Standard(model), [Planar(2), Radial(2), Planar(2), Radial(2), Planar(2)], model
+        Standard(model), [FullAffine(2), Planar(2)], model
     )
 
     # Optimize variational distribution
@@ -126,7 +133,7 @@ def test_flows(benchmark, var_draws):
         vari.fit(10000, var_draws=var_draws)
 
     benchmark(benchmark_fit)
-    vari = vari.fit(10000)
+    vari = vari.fit(100)
 
     mean = vari.sample(1000).mean(0)
     var = vari.sample(1000).var(0)