Update naming of train scripts

docs/examples/conditional.ipynb

@@ -35,7 +35,7 @@
     "import numpy as onp\n",
     "from flowjax.flows import BlockNeuralAutoregressiveFlow\n",
     "from flowjax.distributions import Normal\n",
-    "from flowjax.train import train_flow\n",
+    "from flowjax.train import fit_to_data\n",
     "import matplotlib.pyplot as plt"
    ]
   },
@@ -83,7 +83,7 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      " 88%|████████▊ | 44/50 [00:07<00:01,  5.65it/s, train=1.27, val=1.47 (Max patience reached)]\n"
+      " 88%|████████▊ | 44/50 [00:09<00:01,  4.85it/s, train=1.27, val=1.47 (Max patience reached)]\n"
      ]
     }
    ],
@@ -97,7 +97,7 @@
     ")\n",
     "\n",
     "key, subkey = jr.split(key)\n",
-    "flow, losses = train_flow(subkey, flow, x, u, learning_rate=1e-2, max_patience=10)"
+    "flow, losses = fit_to_data(subkey, flow, x, u, learning_rate=1e-2, max_patience=10)"
    ]
   },
   {

docs/examples/unconditional.ipynb

@@ -28,7 +28,7 @@
     "import jax.numpy as jnp\n",
     "import jax.random as jr\n",
     "from flowjax.flows import MaskedAutoregressiveFlow\n",
-    "from flowjax.train import train_flow\n",
+    "from flowjax.train import fit_to_data\n",
     "from flowjax.distributions import Normal\n",
     "from flowjax.bijections import RationalQuadraticSpline\n",
     "import matplotlib.pyplot as plt\n",
@@ -109,7 +109,7 @@
    ],
    "source": [
     "key, subkey = jr.split(key)\n",
-    "flow, losses = train_flow(subkey, flow, x, learning_rate=1e-3)"
+    "flow, losses = fit_to_data(subkey, flow, x, learning_rate=1e-3)"
    ]
   },
   {
@@ -128,7 +128,7 @@
     {
      "data": {
       "text/plain": [
-       "DeviceArray([-6.6872644, -6.6872644, -6.6872644, -6.6872644, -6.6872644],            dtype=float32)"
+       "Array([-6.6872644, -6.6872644, -6.6872644, -6.6872644, -6.6872644],      dtype=float32)"
       ]
      },
      "execution_count": 5,

flowjax/train/__init__.py

@@ -1,2 +1,2 @@
-from flowjax.train.data_fit import train_flow
-from flowjax.train.variational_fit import variational_fit
+from flowjax.train.data_fit import fit_to_data
+from flowjax.train.variational_fit import fit_to_variational_target

flowjax/train/data_fit.py

@@ -13,7 +13,7 @@
 from flowjax.utils import Array
 
 
-def train_flow(
+def fit_to_data(
     key: KeyArray,
     dist: Distribution,
     x: Array,

flowjax/train/variational_fit.py

@@ -21,19 +21,19 @@
 
 @eqx.filter_jit
 def elbo_loss(dist: Distribution, target: VariationalTarget, key: random.KeyArray, elbo_samples: int = 500):
-    samples = dist.sample(key, sample_shape=(elbo_samples,))
-    approx_density = dist.log_prob(samples).reshape(-1)
-    target_density = target(samples).reshape(-1)
+    samples, approx_density = dist.sample_and_log_prob(key, sample_shape=(elbo_samples,))
+    target_density = target(samples)
     losses = approx_density - target_density
     return losses.mean()
 
-def variational_fit(
+def fit_to_variational_target(
     key: random.KeyArray,
     dist: Distribution,
     target: VariationalTarget,
     loss_fn: VariationalLoss = elbo_loss,
+    learning_rate: float = 5e-4,
+    clip_norm: float = 0.5,
     num_epochs: int = 100,
-    optimizer: Optional[optax.GradientTransformation] = None,
     show_progress: bool = True,
 ):
     """
@@ -44,9 +44,9 @@ def variational_fit(
         dist (Distribution): Distribution object, trainable parameters are found using equinox.is_inexact_array.
         target (VariationalTarget): The target (usually) unormalized log posterior.
         loss_fcn (VariationalLoss, optional): Loss function. Defaults to elbo_loss.
+        learning_rate (float, optional): Adam learning rate. Defaults to 5e-4.
+        clip_norm (float, optional): Maximum gradient norm before clipping occurs. Defaults to 0.5.
         num_epochs (int, optional): The number of training steps to run. Defaults to 100.
-        optimizer (Optional[optax.Optimizer], optional): An optax optimizer (optimizers are implemented as GradientTransformation objects). 
-                                                         Defaults to an adam optimizer with learning rate 5e-4.
         show_progress (bool, optional): Whether to show progress bar. Defaults to True.
     """
     @eqx.filter_jit
@@ -56,12 +56,10 @@ def step(dist, target, key, optimizer, opt_state):
         dist = eqx.apply_updates(dist, updates)
         return dist, opt_state, loss_val
 
-    # Set up a default optimizer if None is provided
-    if optimizer is None:
-        optimizer = optax.chain(
-            optax.clip_by_global_norm(0.5), 
-            optax.adam(learning_rate=5e-4)
-        )
+    optimizer = optax.chain(
+        optax.clip_by_global_norm(clip_norm), 
+        optax.adam(learning_rate=learning_rate)
+    )
 
     trainable_params, _ = eqx.partition(dist, eqx.is_inexact_array)
     opt_state = optimizer.init(trainable_params)
@@ -78,4 +76,4 @@ def step(dist, target, key, optimizer, opt_state):
         if show_progress:
             loop.set_postfix({'loss': losses[-1]})
 
-    return dist, losses
+    return dist, losses, optimizer

tests/train/test_train_utils.py → tests/train/test_data_fit.py

@@ -8,7 +8,7 @@
 from flowjax.distributions import Normal, Transformed
 from flowjax.train.data_fit import (
     count_fruitless,
-    train_flow,
+    fit_to_data,
     train_val_split,
 )
 
@@ -40,7 +40,7 @@ def test_train_flow_filter_spec():
     # All params should change by default
     before = eqx.filter(flow, eqx.is_inexact_array)
     x = random.normal(random.PRNGKey(0), (100, dim))
-    flow, _ = train_flow(random.PRNGKey(0), flow, x, max_epochs=1, batch_size=50)
+    flow, _ = fit_to_data(random.PRNGKey(0), flow, 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)
@@ -50,7 +50,7 @@ def test_train_flow_filter_spec():
     before = eqx.filter(flow, eqx.is_inexact_array)
     filter_spec = jtu.tree_map(lambda x: eqx.is_inexact_array(x), flow)
     filter_spec = eqx.tree_at(lambda tree: tree.base_dist, filter_spec, replace=False)
-    flow, _ = train_flow(
+    flow, _ = fit_to_data(
         random.PRNGKey(0), flow, x, max_epochs=1, batch_size=50, filter_spec=filter_spec
     )
     after = eqx.filter(flow, eqx.is_inexact_array)