Small fixes

thebuckleylab · Jun 28, 2024 · f8ff3a3 · f8ff3a3
1 parent 371adc9
commit f8ff3a3
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/jpc/_core/_energies.py b/jpc/_core/_energies.py
@@ -27,15 +27,15 @@ def pc_energy_fn(
 
     !!! note
 
-        The input and output correspond to the prior and observation of
+        The input x and output y correspond to the prior and observation of
         the generative model, respectively.
 
     **Main arguments:**
 
     - `model`: List of callable model (e.g. neural network) layers.
     - `activities`: List of activities for each layer free to vary.
-    - `output`: Observation or target of the generative model.
-    - `input`: Optional prior of the generative model.
+    - `y`: Observation or target of the generative model.
+    - `x`: Optional prior of the generative model.
 
     **Other arguments:**
 

diff --git a/jpc/_utils.py b/jpc/_utils.py
@@ -119,6 +119,7 @@ def loop_body(state):
         energies_iters = energies_iters.at[:, t].set(energies)
         return t + 1, energies_iters
 
+    # 4096 is the max number of steps set in diffrax
     energies_iters = zeros((len(model), 4096))
     _, energies_iters = jax.lax.while_loop(
         lambda state: state[0] < t_max,