Add documentation for experimental compile mode.

PiperOrigin-RevId: 794238903

Author: Nush395

docs/running.rst

@@ -77,6 +77,18 @@ internal TORAX functions. Used for debugging.
 
   export TORAX_COMPILATION_ENABLED=<True/False>
 
+``EXPERIMENTAL_COMPILE`` (default: `False`) - If `True`, trigger expanded
+compilation up to the TORAX step function's call method. Normally with
+`TORAX_COMPILATION_ENABLED` flag set to `True` TORAX will compile a more limited
+set of functions. This is an experimental flag at the moment and will lead to
+faster simulations (at the expense of larger compilation time). We are currently
+working on reducing compile times so that the behaviour of this flag is
+triggered by `TORAX_COMPILATION_ENABLED`.
+
+.. code-block:: console
+
+  export EXPERIMENTAL_COMPILE=<True/False>
+
 .. _torax_flags:
 
 run_torax flags

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 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_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
+
+  (final_state, _), _ = jax.lax.scan(
+      scan_body, initial_scan_carry, length=max_steps, unroll=scan_unroll
+  )
+
+  return final_state

torax/_src/solver/pydantic_model.py

@@ -55,7 +55,9 @@ class BaseSolver(torax_pydantic.BaseModelFrozen, abc.ABC):
       torax_pydantic.UnitInterval, torax_pydantic.JAX_STATIC
   ] = 1.0
   use_predictor_corrector: Annotated[bool, torax_pydantic.JAX_STATIC] = False
-  n_corrector_steps: pydantic.PositiveInt = 10
+  n_corrector_steps: Annotated[
+      pydantic.PositiveInt, torax_pydantic.JAX_STATIC
+  ] = 10
   convection_dirichlet_mode: Annotated[
       Literal['ghost', 'direct', 'semi-implicit'], torax_pydantic.JAX_STATIC
   ] = 'ghost'

torax/_src/solver/runtime_params.py

@@ -36,4 +36,4 @@ class DynamicRuntimeParams:
 
   chi_pereverzev: float
   D_pereverzev: float  # pylint: disable=invalid-name
-  n_corrector_steps: int
+  n_corrector_steps: int = dataclasses.field(metadata={'static': True})

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=4):
+      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()