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: 794095925

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,46 @@ 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'])
+def max_steps_while_loop(
+    cond_fun: Callable[[_State], BooleanNumeric],
+    body_fun: Callable[[_State], _State],
+    init_val: _State,
+    max_steps: int,
+) -> _State:
+  """A reverse-mode differentiable while_loop using jax.lax.scan.
+
+  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.
+
+  Returns:
+    The final state after the loop terminates 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_met_prev = carry
+
+    cond_eval = cond_fun(current_state)
+    should_execute_body = jnp.logical_and(cond_met_prev, cond_eval)
+
+    next_state = jax.lax.cond(
+        should_execute_body, body_fun, lambda s: s, current_state
+    )
+    next_cond_met = should_execute_body
+
+    return (next_state, next_cond_met), None
+
+  dummy_xs = jnp.arange(max_steps)
+
+  (final_state, _), _ = jax.lax.scan(scan_body, initial_scan_carry, dummy_xs)
+
+  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.max_steps_while_loop(
+          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=4):
+      init_state = (0, x)
+      return jax_utils.max_steps_while_loop(
+          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.max_steps_while_loop(
+          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()