Add reverse-mode differentiable jax.lax.while_loop to jax_utils.

This will be useful for being able to reverse mode differentiate through the TORAX simulation.

PiperOrigin-RevId: 796482656

Author: Nush395

torax/_src/jax_utils.py

@@ -18,7 +18,7 @@
 import functools
 import inspect
 import os
-from typing import Any, Callable, Literal, TypeVar
+from typing import Any, Callable, Literal, ParamSpec, TypeAlias, TypeVar
 
 import chex
 import equinox as eqx
@@ -27,7 +27,8 @@
 import numpy as np
 
 T = TypeVar('T')
-BooleanNumeric = Any  # A bool, or a Boolean array.
+BooleanNumeric: TypeAlias = Any  # A bool, or a Boolean array.
+_State = ParamSpec('_State')
 
 
 @functools.cache
@@ -296,3 +297,55 @@ def init_array(x):
       return x
 
   return jax.tree_util.tree_map(init_array, t)
+
+
+@functools.partial(
+    jit, static_argnames=['cond_fun', 'body_fun', 'max_steps', 'scan_unroll']
+)
+def while_loop_bounded(
+    cond_fun: Callable[[_State], BooleanNumeric],
+    body_fun: Callable[[_State], _State],
+    init_val: _State,
+    max_steps: int,
+    scan_unroll: int = 1,
+) -> _State:
+  """A reverse-mode differentiable while_loop.
+
+  This makes use of jax.lax.scan and `max_steps` to define a fixed size
+  computational graph. The body_fun is called the same number of times it would
+  be under a jax.lax.while_loop i.e. until `cond_fun` returns False (unless the
+  `max_steps` is reached).
+
+  Args:
+    cond_fun: As in jax.lax.while_loop.
+    body_fun: As in jax.lax.while_loop.
+    init_val: As in jax.lax.while_loop.
+    max_steps: An integer, the maximum number of iterations the loop can
+      perform. This is crucial for defining a fixed computational graph for
+      scan.
+    scan_unroll: The number of iterations to unroll the internal scan by.
+
+  Returns:
+    The final state after `cond_fun` returns `False` or `max_steps` are reached.
+  """
+  # Initial carry for the scan: (current_state, while_loop_condition_met)
+  initial_scan_carry = (init_val, jnp.array(True, dtype=jnp.bool_))
+
+  def scan_body(carry, _):
+    current_state, cond_prev = carry
+    # Only execute cond if the previous cond was True.
+    should_execute_body = jax.lax.cond(
+        cond_prev, cond_fun, lambda _: False, current_state
+    )
+    # If the `while_loop` would have terminated, we no-op.
+    next_state = jax.lax.cond(
+        should_execute_body, body_fun, lambda s: s, current_state
+    )
+
+    return (next_state, should_execute_body), None
+
+  (final_state, _), _ = jax.lax.scan(
+      scan_body, initial_scan_carry, length=max_steps, unroll=scan_unroll
+  )
+
+  return final_state

torax/_src/tests/jax_utils_test.py

@@ -150,6 +150,58 @@ def f(x, z, y=2.0):
     x = {'temp1': jnp.array(1.3), 'temp2': jnp.array(2.6)}
     chex.assert_trees_all_close(f_non_inlined(x, z='left'), f(x, z='left'))
 
+  def test_max_steps_while_loop(self):
+    terminating_step = 4
+
+    def cond_fun(state):
+      i, _ = state
+      return i < terminating_step
+
+    def body_fun(state):
+      i, value = state
+      next_i = i + 1
+      next_value = jnp.sin(value)
+      return next_i, next_value
+
+    init_state = (0, 0.5)
+    max_steps = 10
+
+    with self.subTest('forward_agrees_with_while_loop'):
+      output_state = jax_utils.while_loop_bounded(
+          cond_fun, body_fun, init_state, max_steps
+      )
+      chex.assert_trees_all_close(
+          output_state, jax.lax.while_loop(cond_fun, body_fun, init_state)
+      )
+
+    def f_while(x, max_steps=max_steps):
+      init_state = (0, x)
+      return jax_utils.while_loop_bounded(
+          cond_fun, body_fun, init_state, max_steps=max_steps
+      )[1]
+
+    def f(x, n_times=terminating_step):
+      result = x
+      for _ in range(n_times):
+        result = jnp.sin(result)
+      return result
+
+    with self.subTest('forward_agrees_with_explicit'):
+      chex.assert_trees_all_close(f_while(0.5), f(0.5))
+    with self.subTest('grad_agrees_with_explicit'):
+      chex.assert_trees_all_close(jax.grad(f_while)(0.5), jax.grad(f)(0.5))
+
+    with self.subTest('max_steps_is_respected'):
+      final_i, final_value = jax_utils.while_loop_bounded(
+          cond_fun, body_fun, init_state, max_steps=2
+      )
+      self.assertEqual(final_i, 2)
+      chex.assert_trees_all_close(final_value, f(0.5, n_times=2))
+      chex.assert_trees_all_close(
+          jax.grad(f_while)(0.5, max_steps=2), jax.grad(f)(0.5, n_times=2)
+      )
+      chex.assert_trees_all_close(f_while(0.5, max_steps=3), f(0.5, n_times=3))
+
 
 if __name__ == '__main__':
   absltest.main()