Make InterpolatedVarSingleAxis a PyTree.

This is to prepare for making the `GeometryProvider`s PyTrees.

Moved the sorting checks in the two classes to the beginning of the container and make JIT compatible.

Speed check on iterhybrid_rampup has no timing change with/without TORAX_ERRORS_ENABLED.

PiperOrigin-RevId: 775675478

Author: Nush395

.github/workflows/pytest.yml

@@ -73,6 +73,8 @@ jobs:
       run: |
         pytest \
           -vv -n auto \
+          --ignore=torax/tests/sim_experimental_compile_test.py \
+          --ignore=torax/tests/sim_no_compile_test.py \
           --shard-id=$((${{ matrix.shard-id }} - 1)) --num-shards=${{ env.PYTEST_NUM_SHARDS }}
     # Two test require an extra environment variable, so we run them separately.
     - name: Run sim_experimental_compile

torax/_src/interpolated_param.py

@@ -126,7 +126,7 @@ def get_value(self, x: chex.Numeric) -> chex.Array:
     """Returns a value for this parameter interpolated at the given input."""
 
 
-class PiecewiseLinearInterpolatedParam(InterpolatedParamBase):
+class _PiecewiseLinearInterpolatedParam(InterpolatedParamBase):
   """Parameter using piecewise-linear interpolation to compute its value."""
 
   def __init__(self, xs: chex.Array, ys: chex.Array):
@@ -149,10 +149,6 @@ def __init__(self, xs: chex.Array, ys: chex.Array):
     if ys.ndim not in (1, 2):
       raise ValueError(f'ys must be either 1D or 2D. Given: {self.ys.shape}.')
 
-    xs_np = np.array(self.xs)
-    if not np.array_equal(np.sort(xs_np), xs_np):
-      raise RuntimeError('xs must be sorted.')
-
   @property
   def xs(self) -> chex.Array:
     return self._xs
@@ -196,7 +192,7 @@ def get_value(
         raise ValueError(f'ys must be either 1D or 2D. Given: {self.ys.shape}.')
 
 
-class StepInterpolatedParam(InterpolatedParamBase):
+class _StepInterpolatedParam(InterpolatedParamBase):
   """Parameter using step interpolation to compute its value."""
 
   def __init__(self, xs: chex.Array, ys: chex.Array):
@@ -211,8 +207,6 @@ def __init__(self, xs: chex.Array, ys: chex.Array):
           'xs and ys must have the same number of elements in the first '
           f'dimension. Given: {self.xs.shape} and {self.ys.shape}.'
       )
-    if not jnp.allclose(jnp.sort(self.xs), self.xs):
-      raise RuntimeError('xs must be sorted.')
 
   @property
   def xs(self) -> chex.Array:
@@ -325,6 +319,7 @@ def convert_input_to_xs_ys(
     )
 
 
+@jax.tree_util.register_pytree_node_class
 class InterpolatedVarSingleAxis(InterpolatedParamBase):
   """Parameter that may vary based on an input coordinate.
 
@@ -364,8 +359,12 @@ def __init__(
       interpolation_mode: Defines how to interpolate between values in `value`.
       is_bool_param: If True, the input value is assumed to be a bool and is
         converted to a float.
+    Raises:
+      RuntimeError: If the input xs is not sorted.
     """
+    self._value = value
     xs, ys = value
+    jax_utils.error_if(xs, jnp.any(jnp.diff(xs) < 0), 'xs must be sorted.')
 
     if not np.issubdtype(xs.dtype, np.floating):
       raise ValueError(f'xs must be a float array, but got {xs.dtype}.')
@@ -376,12 +375,23 @@ def __init__(
     self._interpolation_mode = interpolation_mode
     match interpolation_mode:
       case InterpolationMode.PIECEWISE_LINEAR:
-        self._param = PiecewiseLinearInterpolatedParam(xs=xs, ys=ys)
+        self._param = _PiecewiseLinearInterpolatedParam(xs=xs, ys=ys)
       case InterpolationMode.STEP:
-        self._param = StepInterpolatedParam(xs=xs, ys=ys)
+        self._param = _StepInterpolatedParam(xs=xs, ys=ys)
       case _:
         raise ValueError('Unknown interpolation mode.')
 
+  def tree_flatten(self):
+    static_params = {
+        'interpolation_mode': self.interpolation_mode,
+        'is_bool_param': self.is_bool_param,
+    }
+    return (self._value, static_params)
+
+  @classmethod
+  def tree_unflatten(cls, aux_data, children):
+    return cls(children, **aux_data)
+
   @property
   def is_bool_param(self) -> bool:
     """Returns whether this param represents a bool."""

torax/_src/tests/interpolated_param_test.py

@@ -16,6 +16,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
@@ -173,8 +174,8 @@ def test_multi_value_range_returns_expected_output(
     )
 
   @parameterized.parameters(
-      (interpolated_param.PiecewiseLinearInterpolatedParam,),
-      (interpolated_param.StepInterpolatedParam,),
+      (interpolated_param._PiecewiseLinearInterpolatedParam,),
+      (interpolated_param._StepInterpolatedParam,),
   )
   def test_interpolated_param_1d_xs_and_1d_or_2d_ys(self, range_class):
     """Tests that the interpolated_param only take 1D inputs."""
@@ -198,8 +199,8 @@ def test_interpolated_param_1d_xs_and_1d_or_2d_ys(self, range_class):
       )
 
   @parameterized.parameters(
-      (interpolated_param.PiecewiseLinearInterpolatedParam,),
-      (interpolated_param.StepInterpolatedParam,),
+      (interpolated_param._PiecewiseLinearInterpolatedParam,),
+      (interpolated_param._StepInterpolatedParam,),
   )
   def test_interpolated_param_need_xs_ys_same_shape(self, range_class):
     """Tests the xs and ys inputs have to have the same shape."""
@@ -214,19 +215,20 @@ def test_interpolated_param_need_xs_ys_same_shape(self, range_class):
       )
 
   @parameterized.parameters(
-      (interpolated_param.PiecewiseLinearInterpolatedParam,),
-      (interpolated_param.StepInterpolatedParam,),
+      interpolated_param.InterpolationMode.PIECEWISE_LINEAR,
+      interpolated_param.InterpolationMode.STEP,
   )
-  def test_interpolated_param_need_xs_to_be_sorted(self, range_class):
-    """Tests the xs inputs have to be sorted."""
-    range_class(
-        xs=jnp.array([1.0, 2.0, 3.0, 4.0]),
-        ys=jnp.array([1.0, 2.0, 3.0, 4.0]),
-    )
+  def test_interpolated_param_is_invariant_to_xs_order(
+      self,
+      interpolation_mode: interpolated_param.InterpolationMode,
+  ):
     with self.assertRaises(RuntimeError):
-      range_class(
-          xs=jnp.array([4.0, 2.0, 1.0, 3.0]),
-          ys=jnp.array([1.0, 2.0, 3.0, 4.0]),
+      interpolated_param.InterpolatedVarSingleAxis(
+          value=(
+              jnp.array([4.0, 2.0, 1.0, 3.0]),
+              jnp.array([4.0, 2.0, 1.0, 3.0]),
+          ),
+          interpolation_mode=interpolation_mode,
       )
 
   @parameterized.named_parameters(
@@ -493,19 +495,23 @@ def test_interpolated_param_get_value_is_jittable(
           interpolated_param.InterpolationMode.STEP,
       ],
   )
-  def test_interpolated_var_properties(
+  def test_interpolated_param_is_usable_under_jit(
       self,
       is_bool: bool,
       interpolation_mode: interpolated_param.InterpolationMode,
   ):
-    """Check the properties of the interpolated var are set correctly."""
     var = interpolated_param.InterpolatedVarSingleAxis(
         value=(np.array([0.0, 1.0]), np.array([0.0, 1.0])),
         is_bool_param=is_bool,
         interpolation_mode=interpolation_mode,
     )
-    self.assertEqual(var.is_bool_param, is_bool)
-    self.assertEqual(var.interpolation_mode, interpolation_mode)
+
+    def f(x: interpolated_param.InterpolatedVarSingleAxis, t: chex.Numeric):
+      return x.get_value(x=t)
+
+    interpolated_output_jit = jax.jit(f)(var, 0.5)
+    interpolated_output = f(var, 0.5)
+    np.testing.assert_allclose(interpolated_output_jit, interpolated_output)
 
 
 if __name__ == '__main__':