Replace functools.cached_property with property in CellVariable.

Using functools.cached_property can cause issues with JAX transformations like jax.lax.cond, as it modifies the object's internal state upon first access, potentially leading to UnexpectedTracerError. Switching to standard properties avoids these side effects within JAX contexts.

This unfortunately means that we won't be able to benefit from using `cached_properties` in CellVariable. The operations to get these properties are simple, so the impact to performance should be negligible.

I added a unittest that exposed the tracer leak to prevent regression in the future.

PiperOrigin-RevId: 861257087

Author: hamelphi

torax/_src/fvm/cell_variable.py

@@ -19,7 +19,6 @@
 [https://www.ctcms.nist.gov/fipy/]
 """
 import dataclasses
-import functools
 
 import chex
 import jax
@@ -131,7 +130,7 @@ class CellVariable:
   # Can't make the above default values be jax zeros because that would be a
   # call to jax before absl.app.run
 
-  @functools.cached_property
+  @property
   def cell_centers(self) -> jt.Float[chex.Array, 'cell']:
     """Locations of the cell centers."""
     return (self.face_centers[..., 1:] + self.face_centers[..., :-1]) / 2.0
@@ -281,7 +280,7 @@ def left_face_value(self) -> jt.Float[chex.Array, '']:
       value = self.value[..., 0:1]
     return value
 
-  @functools.cached_property
+  @property
   def right_face_value(self) -> jt.Float[chex.Array, '']:
     """Calculates the value of the rightmost face."""
     if self.right_face_constraint is not None:

torax/_src/fvm/tests/cell_variable_test.py

@@ -14,6 +14,7 @@
 from absl.testing import absltest
 from absl.testing import parameterized
 import chex
+import jax
 from jax import numpy as jnp
 import numpy as np
 from torax._src.fvm import cell_variable
@@ -441,6 +442,43 @@ def test_compute_face_grad_compared_to_forward_diff(self):
         accurate_method_error, forward_difference_error
     )
 
+  def test_jax_tracer_leak_in_cond(self):
+    """Tests that accessing properties inside jax.cond does not leak tracers."""
+    # This test verifies that we don't have side effects (like cached_property
+    # updates) that leak tracers from inside a jax.cond to the outer scope.
+
+    n = 10
+    value = jnp.zeros((n,))
+    face_centers = jnp.linspace(0, 1, n + 1)
+
+    # Create the object OUTSIDE the cond
+    var = cell_variable.CellVariable(value=value, face_centers=face_centers)
+
+    def true_branch(_):
+      # Access properties INSIDE the cond
+      # If these are cached_property and write to var.__dict__, it might leak
+      _ = var.cell_centers
+      _ = var.right_face_value
+      return var.cell_centers
+
+    def false_branch(_):
+      return jnp.zeros((n,))
+
+    jax.lax.cond(jnp.array(True), true_branch, false_branch, None)
+
+    # Access the property AFTER the cond to ensure no leaked tracer is stored.
+    centers = var.cell_centers
+    right_face = var.right_face_value
+
+    # Force use of the values in a way that triggers UnexpectedTracerError for
+    # leaked tracers
+    val_check = centers + 1.0
+    right_check = right_face + 1.0
+
+    # Casting to numpy also forces evaluation
+    np.array(val_check)
+    np.array(right_check)
+
 
 if __name__ == '__main__':
   absltest.main()