Add a BatchedCellVariable class.

This will allow for vectorizing code that currently loops across sequences of CellVariables.

PiperOrigin-RevId: 795051091

Author: sbodenstein

torax/_src/fvm/cell_variable.py

@@ -20,18 +20,20 @@
 [https://www.ctcms.nist.gov/fipy/]
 """
 import dataclasses
+from typing import Self, Sequence
 
 import chex
 import jax
 from jax import numpy as jnp
 import jaxtyping as jt
 from torax._src import array_typing
+from torax._src import jax_utils
 import typing_extensions
 
 
 def _zero() -> array_typing.FloatScalar:
   """Returns a scalar zero as a jax Array."""
-  return jnp.zeros(())
+  return jnp.zeros((), dtype=jax_utils.get_dtype())
 
 
 @chex.dataclass(frozen=True)
@@ -266,3 +268,62 @@ def __eq__(self, other: typing_extensions.Self) -> bool:
       return True
     except AssertionError:
       return False
+
+
+@dataclasses.dataclass(frozen=True)
+class BatchedCellVariable:
+  """An object representing a batch of `CellVariable` objects."""
+
+  value: jt.Float[jax.Array, 'batch cell']
+  dr: jt.Float[jax.Array, 'batch']
+  left_face_constraint: jt.Float[jax.Array, 'batch']
+  right_face_constraint: jt.Float[jax.Array, 'batch']
+  left_face_grad_constraint: jt.Float[jax.Array, 'batch']
+  right_face_grad_constraint: jt.Float[jax.Array, 'batch']
+
+  @classmethod
+  def construct(cls, x: CellVariable | Sequence[CellVariable]) -> Self:
+    """Constructs a `CellVariableBatched` object.
+
+    Note that any `None` values are represented as `jnp.nan` values.
+
+    Args:
+      x: A `CellVariable` or a sequence of `CellVariables`.
+
+    Returns:
+      A `BatchedCellVariable` object.
+    """
+    x = (x,) if isinstance(x, CellVariable) else x
+    if not x:
+      raise ValueError('The tuple must have at least one element.')
+
+    x = [_conform(v) for v in x]
+
+    def concat(*x):
+      return jnp.concatenate(x, axis=0)
+
+    x = jax.tree_util.tree_map(concat, *x)
+    return cls(**x)
+
+
+def _conform(x: CellVariable):
+  def conform_scalar(x):
+    if x is None:
+      return jnp.array([jnp.nan], dtype=jax_utils.get_dtype())
+
+    if x.ndim == 0:
+      return jnp.expand_dims(x, axis=0)
+    return x
+
+  value = jnp.expand_dims(x.value, axis=0) if x.value.ndim == 1 else x.value
+
+  return {
+      'value': value,
+      'dr': conform_scalar(x.dr),
+      'left_face_constraint': conform_scalar(x.left_face_constraint),
+      'right_face_constraint': conform_scalar(x.right_face_constraint),
+      'left_face_grad_constraint': conform_scalar(x.left_face_grad_constraint),
+      'right_face_grad_constraint': conform_scalar(
+          x.right_face_grad_constraint
+      ),
+  }

torax/_src/fvm/tests/cell_variable_test.py

@@ -584,6 +584,44 @@ def test_almost_equal(
       with self.assertRaises(AssertionError):
         chex.assert_trees_all_close(var1, var2, atol=atol)
 
+  def test_construct_batched_cell_variable(self):
+    var_1 = cell_variable.CellVariable(
+        value=jnp.array([1.0, 2.0]),
+        dr=jnp.array(0.1),
+        left_face_constraint=jnp.array(3.0),
+        left_face_grad_constraint=None,
+    )
+    var_2 = cell_variable.CellVariable(
+        value=jnp.array([5.0, 6.0]),
+        dr=jnp.array(0.2),
+        right_face_constraint=jnp.array(5.0),
+        right_face_grad_constraint=None,
+    )
+    with self.subTest('from_tuple'):
+      var = cell_variable.BatchedCellVariable.construct((var_1, var_2))
+      chex.assert_trees_all_equal(
+          var.value, jnp.array([[1.0, 2.0], [5.0, 6.0]])
+      )
+      chex.assert_trees_all_equal(var.dr, jnp.array([0.1, 0.2]))
+
+      chex.assert_trees_all_equal(
+          var.left_face_constraint, jnp.array([3.0, jnp.nan])
+      )
+      chex.assert_trees_all_equal(
+          var.right_face_constraint, jnp.array([jnp.nan, 5.0])
+      )
+      chex.assert_trees_all_equal(
+          var.left_face_grad_constraint, jnp.array([jnp.nan, 0.0])
+      )
+      chex.assert_trees_all_equal(
+          var.right_face_grad_constraint, jnp.array([0.0, jnp.nan])
+      )
+
+    with self.subTest('from_single_var'):
+      var = cell_variable.BatchedCellVariable.construct(var_1)
+      chex.assert_trees_all_equal(var.value, jnp.array([[1.0, 2.0]]))
+      chex.assert_trees_all_equal(var.dr, jnp.array([0.1]))
+
 
 if __name__ == '__main__':
   absltest.main()

torax/_src/jax_utils.py

@@ -29,6 +29,7 @@
 T = TypeVar('T')
 BooleanNumeric: TypeAlias = Any  # A bool, or a Boolean array.
 _State = ParamSpec('_State')
+PyTree: TypeAlias = Any
 
 
 @functools.cache
@@ -299,6 +300,70 @@ def init_array(x):
   return jax.tree_util.tree_map(init_array, t)
 
 
+def batched_cond(
+    pred: jax.Array,
+    true_fun: Callable[..., PyTree],
+    false_fun: Callable[..., PyTree],
+    operands: tuple[PyTree, ...],
+    implementation: Literal['vectorize', 'map'] = 'vectorize',
+):
+  """A batched version of `jax.lax.cond`.
+
+  JAX provides two approaches for implementing a batched version of
+  `jax.lax.cond`, neither of which is always faster:
+  `implementation='vectorize'` is equivalent to `jnp.select`, which evaluates
+  both braches for every batch element. This is fully vectorized, allowing for
+  parallel execution on CPU/GPU, but requiring twice the number of function
+  evaluations. `implementation='map'` will sequentially evaluate `jax.lax.cond`,
+  preventing vectorized execution, but only requiring a single function
+  evaluation per batch element.
+
+  This function also handles the special case where `pred` is a concrete list of
+  length-1, in which case we can avoid tracing both branches like `jax.lax.cond`
+  does by doing the control-flow in Python.
+
+  Args:
+    pred: Boolean 1D array `[batch_size]`, indicating which branch function to
+      apply.
+    true_fun: Function (A -> B), to be applied if `pred` is True.
+    false_fun: Function (A -> B), to be applied if `pred` is False.
+    operands: A tuple of arguments to pass to the functions. Each `jax.Array`
+      (every PyTree leaf) must have a leading batch dimension of size
+      `batch_size`.
+    implementation: The implementation to use. 'vectorize' compiles to a
+      `jax.lax.select`, where both branches are evaluated. 'map' uses
+      `jax.lax.map`.
+
+  Returns:
+    The result of applying the appropriate function to each element of the
+    batch.
+  """
+
+  if not isinstance(operands, tuple):
+    raise ValueError('The args must be a tuple.')
+
+  if pred.ndim != 1 or pred.dtype != jnp.bool:
+    raise ValueError('pred must be a 1D array of bools.')
+
+  # For the special case where `pred` is a concrete list of length 1, we can
+  # avoid tracing both branches by doing the control flow in Python.
+  if len(pred) == 1 and not isinstance(pred, jax.core.Tracer):
+    f = true_fun if bool(pred) else false_fun
+    operands = jax.tree.map(lambda x: jnp.squeeze(x, axis=0), operands)
+    out = f(*operands)
+    return jax.tree.map(lambda x: jnp.expand_dims(x, axis=0), out)
+
+  f = lambda args: jax.lax.cond(args[0], true_fun, false_fun, *args[1])
+  match implementation:
+    case 'vectorize':
+      # This is compiled to a jax.lax.select, where both branches are evaluated.
+      return jax.vmap(f)((pred, operands))
+    case 'map':
+      return jax.lax.map(f, (pred, operands))
+    case _:
+      raise ValueError(f'Unknown implementation: {implementation}')
+
+
 @functools.partial(
     jit, static_argnames=['cond_fun', 'body_fun', 'max_steps', 'scan_unroll']
 )

torax/_src/tests/jax_utils_test.py

@@ -150,6 +150,41 @@ 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'))
 
+  @parameterized.parameters(['map', 'vectorize'])
+  def test_batched_cond(self, implementation):
+    pred = jnp.array([True, False])
+    x = jnp.array([[2, 3.0, 4.0], [5.0, 6.0, 7.0]])
+    out = jax_utils.batched_cond(
+        pred=pred,
+        true_fun=lambda x, y: x * y,
+        false_fun=lambda x, y: x * y**2,
+        operands=(x, x),
+        implementation=implementation,
+    )
+    out_gt = jnp.array(
+        [[4.0, 9.0, 16.0], [125.0, 216.0, 343.0]], dtype=jnp.float32
+    )
+    chex.assert_trees_all_equal(out, out_gt)
+
+  @parameterized.parameters(['map', 'vectorize'])
+  def test_batched_cond_concrete_special(self, implementation):
+    pred = jnp.array([True])
+    x = jnp.array([[2, 3.0, 4.0]])
+
+    @jax.jit
+    def f(x):
+      return jax_utils.batched_cond(
+          pred=pred,
+          true_fun=lambda x, y: x * y,
+          false_fun=lambda x, y: x * y**2,
+          operands=(x, x),
+          implementation=implementation,
+      )
+
+    out = f(x)
+    out_gt = jnp.array([[4.0, 9.0, 16.0]], dtype=jnp.float32)
+    chex.assert_trees_all_equal(out, out_gt)
+
   def test_max_steps_while_loop(self):
     terminating_step = 4