feat: add unroll param to make_ode and improve docs

- Added `unroll` parameter to `vbjax.make_ode` to match `make_sde` API and enable loop unrolling optimization.
- Updated `make_ode` docstring to include `method` and `unroll` parameters.
- Updated `docs/integrators.rst` to match the implementation signature and document the new parameter.
- Added `test_ode_unroll` test case.

Author: google-labs-jules[bot]

docs/integrators.rst

@@ -120,12 +120,13 @@ make_ode
 
 Create an ordinary differential equation integrator.
 
-.. function:: make_ode(dt, dfun, method='heun', adhoc=None)
+.. function:: make_ode(dt, dfun, adhoc=None, method='heun', unroll=10)
 
    :param float dt: Time step size
    :param callable dfun: Function ``dfun(x, p)`` that returns dx/dt
-   :param str method: Integration method - ``'euler'``, ``'heun'``, or ``'rk4'``
    :param callable adhoc: Optional function ``f(x, p)`` for post-step corrections
+   :param str method: Integration method - ``'euler'``, ``'heun'``, or ``'rk4'``
+   :param int unroll: Loop unroll factor for performance
    :return: Tuple of (step, loop) functions
    :rtype: tuple
 

vbjax/loops.py

@@ -175,7 +175,7 @@ def op(x, z):
     return step, loop
 
 
-def make_ode(dt, dfun, adhoc=None, method='heun'):
+def make_ode(dt, dfun, adhoc=None, method='heun', unroll=10):
     """Use a Heun scheme to integrate autonomous ordinary differential
     equations (ODEs).
 
@@ -189,6 +189,10 @@ def make_ode(dt, dfun, adhoc=None, method='heun'):
     adhoc : function or None
         Function of the form `f(x, p)` that allows making adhoc corrections
         to states after a step.
+    method : str
+        Integration method, one of 'euler', 'heun', 'rk4'.
+    unroll : int
+        Force unrolls the time stepping loop.
 
     Returns
     =======
@@ -226,7 +230,7 @@ def loop(x0, ts, p):
         def op(x, t):
             x = step(x, t, p)
             return x, x
-        return jax.lax.scan(op, x0, ts)[1]
+        return jax.lax.scan(op, x0, ts, unroll=unroll)[1]
 
     return step, loop
 

vbjax/tests/test_loops.py

@@ -22,6 +22,16 @@ def test_ode():
     assert xs.shape == (64, 32)
 
 
+def test_ode_unroll():
+    f = lambda x, _: -x
+    dt = 0.1
+    # Check if unroll parameter is accepted and works
+    _, run = vb.make_ode(dt, f, unroll=5)
+    x0 = np.r_[:32].astype('f')
+    xs = run(x0, np.r_[:64], None)
+    assert xs.shape == (64, 32)
+
+
 def test_dde():
     def dfun(xt, x, t, p):
         xd = xt[0, t-100]