Merge pull request #215 from danielward27/avoid_array_doctest

Avoid comparing array values doctest

Author: danielward27

docs/getting_started.rst

@@ -36,10 +36,12 @@ Additionally, we can evaluate the log probabilities of these samples
 
 .. doctest:: 
 
-   >>> normal.log_prob(sample)
-   Array(-3.4016984, dtype=float32)
-   >>> normal.log_prob(batch)
-   Array([-4.8808994, -5.0121717, -3.2557464, -4.131773 ], dtype=float32)
+   >>> log_prob = normal.log_prob(sample)
+   >>> log_prob.shape
+   ()
+   >>> log_probs = normal.log_prob(batch)
+   >>> log_probs.shape
+   (4,)
 
 When ``sample.shape == distribution.shape``, a scalar log probability is returned. For 
 a batch of samples, the shape of the returned log probabilities matches the shape

flowjax/bijections/affine.py

@@ -180,8 +180,7 @@ class AdditiveCondition(AbstractBijection):
             >>> bijection = AdditiveCondition(
             ...     Linear(2, 3, key=jr.key(0)), shape=(3,), cond_shape=(2,)
             ...     )
-            >>> bijection.transform(jnp.ones(3), condition=jnp.ones(2))
-            Array([1.9670618, 0.8156546, 1.7763454], dtype=float32)
+            >>> y = bijection.transform(jnp.ones(3), condition=jnp.ones(2))
 
     """
 

pyproject.toml

@@ -23,7 +23,7 @@ license = { file = "LICENSE" }
 name = "flowjax"
 readme = "README.md"
 requires-python = ">=3.10"
-version = "17.1.0"
+version = "17.1.1"
 
 [project.urls]
 repository = "https://github.com/danielward27/flowjax"

tests/test_train/test_data_fit.py

@@ -1,11 +1,39 @@
+import equinox as eqx
 import jax.numpy as jnp
 import jax.random as jr
 import pytest
+from jax import random
 
-from flowjax.distributions import Normal, StandardNormal
-from flowjax.train.loops import fit_to_key_based_loss
+from flowjax.bijections import Affine
+from flowjax.distributions import Normal, StandardNormal, Transformed
+from flowjax.train.loops import fit_to_data, fit_to_key_based_loss
 from flowjax.train.losses import ElboLoss
 
+
+def test_data_fit():
+    dim = 3
+    mean, std = jnp.ones(dim), jnp.ones(dim)
+    base_dist = Normal(mean, std)
+    flow = Transformed(base_dist, Affine(jnp.ones(dim), jnp.ones(dim)))
+
+    # All params should change by default
+    before = eqx.filter(flow, eqx.is_inexact_array)
+    x = random.normal(random.key(0), (100, dim))
+    flow, losses = fit_to_data(
+        random.key(0),
+        dist=flow,
+        x=x,
+        max_epochs=1,
+        batch_size=50,
+    )
+    after = eqx.filter(flow, eqx.is_inexact_array)
+
+    assert jnp.all(before.base_dist.bijection.loc != after.base_dist.bijection.loc)
+    assert jnp.all(before.bijection.loc != after.bijection.loc)
+    assert isinstance(losses["train"][0], float)
+    assert isinstance(losses["val"][0], float)
+
+
 test_shapes = [(), (2,), (2, 3, 4)]