Remove batch dimension from CellVariable.

PiperOrigin-RevId: 795051091

Author: sbodenstein

torax/_src/fvm/cell_variable.py

@@ -26,15 +26,18 @@
 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)
+@array_typing.jaxtyped
+@jax.tree_util.register_dataclass
+@dataclasses.dataclass(frozen=True)
 class CellVariable:
   """A variable representing values of the cells along the radius.
 
@@ -53,15 +56,14 @@ class CellVariable:
       of the gradient on the rightmost face variable.
   """
 
-  # t* means match 0 or more leading time dimensions.
-  value: jt.Float[chex.Array, 't* cell']
-  dr: jt.Float[chex.Array, 't*']
-  left_face_constraint: jt.Float[chex.Array, 't*'] | None = None
-  right_face_constraint: jt.Float[chex.Array, 't*'] | None = None
-  left_face_grad_constraint: jt.Float[chex.Array, 't*'] | None = (
+  value: jt.Float[jax.Array, 'cell']
+  dr: array_typing.FloatScalar
+  left_face_constraint: array_typing.FloatScalar | None = None
+  right_face_constraint: array_typing.FloatScalar | None = None
+  left_face_grad_constraint: array_typing.FloatScalar | None = (
       dataclasses.field(default_factory=_zero)
   )
-  right_face_grad_constraint: jt.Float[chex.Array, 't*'] | None = (
+  right_face_grad_constraint: array_typing.FloatScalar | None = (
       dataclasses.field(default_factory=_zero)
   )
   # Can't make the above default values be jax zeros because that would be a
@@ -119,21 +121,9 @@ def __post_init__(self):
           'right_face_grad_constraint must be set.'
       )
 
-  def _assert_unbatched(self):
-    if len(self.value.shape) != 1:
-      raise AssertionError(
-          'CellVariable must be unbatched, but has `value` shape '
-          f'{self.value.shape}. Consider using vmap to batch the function call.'
-      )
-    if self.dr.shape:
-      raise AssertionError(
-          'CellVariable must be unbatched, but has `dr` shape '
-          f'{self.dr.shape}. Consider using vmap to batch the function call.'
-      )
-
   def face_grad(
-      self, x: jt.Float[chex.Array, 'cell'] | None = None
-  ) -> jt.Float[chex.Array, 'face']:
+      self, x: jt.Float[array_typing.Array, 'cell'] | None = None
+  ) -> jt.Float[jax.Array, 'face']:
     """Returns the gradient of this value with respect to the faces.
 
     Implemented using forward differencing of cells. Leftmost and rightmost
@@ -146,7 +136,6 @@ def face_grad(
     Returns:
       A jax.Array of shape (num_faces,) containing the gradient.
     """
-    self._assert_unbatched()
     if x is None:
       forward_difference = jnp.diff(self.value) / self.dr
     else:
@@ -194,7 +183,7 @@ def constrained_grad(
     right = jnp.expand_dims(right_grad, axis=0)
     return jnp.concatenate([left, forward_difference, right])
 
-  def _left_face_value(self) -> jt.Float[chex.Array, '#t']:
+  def _left_face_value(self) -> jt.Float[jax.Array, '#t']:
     """Calculates the value of the leftmost face."""
     if self.left_face_constraint is not None:
       value = self.left_face_constraint
@@ -206,7 +195,7 @@ def _left_face_value(self) -> jt.Float[chex.Array, '#t']:
       value = self.value[..., 0:1]
     return value
 
-  def _right_face_value(self) -> jt.Float[chex.Array, '#t']:
+  def _right_face_value(self) -> jt.Float[jax.Array, '#t']:
     """Calculates the value of the rightmost face."""
     if self.right_face_constraint is not None:
       value = self.right_face_constraint
@@ -222,14 +211,14 @@ def _right_face_value(self) -> jt.Float[chex.Array, '#t']:
       )
     return value
 
-  def face_value(self) -> jt.Float[jax.Array, 't* face']:
+  def face_value(self) -> jt.Float[jax.Array, 'face']:
     """Calculates values of this variable on the face grid."""
     inner = (self.value[..., :-1] + self.value[..., 1:]) / 2.0
     return jnp.concatenate(
         [self._left_face_value(), inner, self._right_face_value()], axis=-1
     )
 
-  def grad(self) -> jt.Float[jax.Array, 't* face']:
+  def grad(self) -> jt.Float[jax.Array, 'face']:
     """Returns the gradient of this variable wrt cell centers."""
     face = self.face_value()
     return jnp.diff(face) / jnp.expand_dims(self.dr, axis=-1)
@@ -251,7 +240,7 @@ def __str__(self) -> str:
     output_string += ')'
     return output_string
 
-  def cell_plus_boundaries(self) -> jt.Float[jax.Array, 't* cell+2']:
+  def cell_plus_boundaries(self) -> jt.Float[jax.Array, 'cell+2']:
     """Returns the value of this variable plus left and right boundaries."""
     right_value = self._right_face_value()
     left_value = self._left_face_value()