Add __hash__ and __eq__ to more objects in anticipation of expanding JIT scope.

PiperOrigin-RevId: 774754949

Author: Nush395

torax/_src/config/build_runtime_params.py

@@ -22,7 +22,11 @@
 `get_consistent_dynamic_runtime_params_slice_and_geometry` which returns a
 DynamicRuntimeParamsSlice and a corresponding geometry with consistent Ip.
 """
+
+import functools
+
 import chex
+import jax
 from torax._src.config import runtime_params_slice
 from torax._src.geometry import geometry
 from torax._src.geometry import geometry_provider as geometry_provider_lib
@@ -89,6 +93,22 @@ def __init__(
     self._mhd = torax_config.mhd
     self._neoclassical = torax_config.neoclassical
 
+  @functools.cached_property
+  def initial_slice(self) -> runtime_params_slice.DynamicRuntimeParamsSlice:
+    return self(0.0)
+
+  @property
+  def structure(self) -> jax.tree_util.PyTreeDef:
+    return jax.tree_util.tree_structure(self.initial_slice)
+
+  def __hash__(self) -> int:
+    """Used for jax.jit caching where we only care about the structure."""
+    return hash(self.structure)
+
+  def __eq__(self, other: typing_extensions.Self) -> bool:
+    """Used for jax.jit caching where we only care about the structure."""
+    return self.structure == other.structure
+
   @classmethod
   def from_config(
       cls,

torax/_src/config/tests/build_runtime_params_test.py

@@ -11,16 +11,18 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
+import os
 from absl.testing import absltest
 from absl.testing import parameterized
 import numpy as np
 from torax._src.config import build_runtime_params
+from torax._src.config import config_loader
 from torax._src.config import profile_conditions as profile_conditions_lib
 from torax._src.geometry import pydantic_model as geometry_pydantic_model
 from torax._src.pedestal_model import pydantic_model as pedestal_pydantic_model
 from torax._src.pedestal_model import set_tped_nped
 from torax._src.test_utils import default_configs
+from torax._src.test_utils import paths
 from torax._src.torax_pydantic import model_config
 from torax._src.torax_pydantic import torax_pydantic
 
@@ -251,6 +253,44 @@ def test_profile_conditions_set_electron_density_and_boundary_condition(
       )
       self.assertTrue(static_slice.n_e_right_bc_is_absolute)
 
+  def test_same_shape_dynamic_params_has_same_hash_and_equals(self):
+    test_data_dir = paths.test_data_dir()
+    torax_config = config_loader.build_torax_config_from_file(
+        os.path.join(test_data_dir, 'test_iterhybrid_rampup.py')
+    )
+    provider1 = (
+        build_runtime_params.DynamicRuntimeParamsSliceProvider.from_config(
+            torax_config
+        )
+    )
+    torax_config.update_fields({'sources.generic_heat.P_total': 1.0})
+    provider2 = (
+        build_runtime_params.DynamicRuntimeParamsSliceProvider.from_config(
+            torax_config
+        )
+    )
+    self.assertEqual(provider1, provider2)
+    self.assertEqual(hash(provider1), hash(provider2))
+
+  def test_different_dynamic_params_has_different_hash_and_not_equals(self):
+    test_data_dir = paths.test_data_dir()
+    torax_config = config_loader.build_torax_config_from_file(
+        os.path.join(test_data_dir, 'test_iterhybrid_rampup.py')
+    )
+    provider1 = (
+        build_runtime_params.DynamicRuntimeParamsSliceProvider.from_config(
+            torax_config
+        )
+    )
+    torax_config.update_fields({'transport': {'model_name': 'constant'}})
+    provider2 = (
+        build_runtime_params.DynamicRuntimeParamsSliceProvider.from_config(
+            torax_config
+        )
+    )
+    self.assertNotEqual(provider1, provider2)
+    self.assertNotEqual(hash(provider1), hash(provider2))
+
 
 if __name__ == '__main__':
   absltest.main()

torax/_src/geometry/geometry_provider.py

@@ -86,6 +86,14 @@ def __call__(
   def torax_mesh(self) -> torax_pydantic.Grid1D:
     """Returns the mesh used by Torax, this is consistent across time."""
 
+  def __hash__(self) -> int:
+    """For the purposes of jax.jit caching, we only care about the mesh."""
+    return hash(self.torax_mesh)
+
+  def __eq__(self, other: typing_extensions.Self) -> bool:
+    """For the purposes of jax.jit caching, we only care about the mesh."""
+    return self.torax_mesh == other.torax_mesh
+
 
 class ConstantGeometryProvider(GeometryProvider):
   """Returns the same Geometry for all calls."""

torax/_src/geometry/tests/geometry_provider_test.py

@@ -74,6 +74,22 @@ def test_none_z_magnetic_axis_stays_none_time_dependent(self):
     self.assertIsNone(provider(0.0)._z_magnetic_axis)
     self.assertIsNone(provider(10.0)._z_magnetic_axis)
 
+  def test_same_mesh_has_same_hash_and_equality(self):
+    geo1 = geometry_pydantic_model.CircularConfig(n_rho=25).build_geometry()
+    geo2 = geometry_pydantic_model.CircularConfig(n_rho=25).build_geometry()
+    provider1 = geometry_provider.ConstantGeometryProvider(geo1)
+    provider2 = geometry_provider.ConstantGeometryProvider(geo2)
+    self.assertEqual(provider1, provider2)
+    self.assertEqual(hash(provider1), hash(provider2))
+
+  def test_different_mesh_has_different_hash_and_not_equals(self):
+    geo1 = geometry_pydantic_model.CircularConfig(n_rho=25).build_geometry()
+    geo2 = geometry_pydantic_model.CircularConfig(n_rho=50).build_geometry()
+    provider1 = geometry_provider.ConstantGeometryProvider(geo1)
+    provider2 = geometry_provider.ConstantGeometryProvider(geo2)
+    self.assertNotEqual(provider1, provider2)
+    self.assertNotEqual(hash(provider1), hash(provider2))
+
 
 if __name__ == "__main__":
   absltest.main()

torax/_src/mhd/base.py

@@ -16,10 +16,17 @@
 
 import chex
 from torax._src.mhd.sawtooth import sawtooth_model
+import typing_extensions
 
 
 @chex.dataclass
 class MHDModels:
   """Container for instantiated MHD model objects."""
 
   sawtooth: sawtooth_model.SawtoothModel | None = None
+
+  def __hash__(self) -> int:
+    return hash(self.sawtooth)
+
+  def __eq__(self, other: typing_extensions.Self) -> bool:
+    return self.sawtooth == other.sawtooth

torax/_src/mhd/tests/base_test.py

@@ -0,0 +1,51 @@
+# Copyright 2024 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+
+from absl.testing import absltest
+from torax._src.config import config_loader
+from torax._src.orchestration import run_simulation
+from torax._src.test_utils import paths
+
+
+class BaseTest(absltest.TestCase):
+
+  def test_different_mhd_models_have_same_hash_and_equals(self):
+    test_data_dir = paths.test_data_dir()
+    torax_config = config_loader.build_torax_config_from_file(
+        os.path.join(test_data_dir, "test_iterhybrid_rampup.py")
+    )
+    static_runtime_params_slice, _, _, _, _, _, step_fn = (
+        run_simulation.prepare_simulation(torax_config)
+    )
+    model1 = torax_config.mhd.build_mhd_models(
+        static_runtime_params_slice,
+        step_fn.solver.transport_model,
+        step_fn.solver.source_models,
+        step_fn.solver.pedestal_model,
+        step_fn.solver.neoclassical_models,
+    )
+    model2 = torax_config.mhd.build_mhd_models(
+        static_runtime_params_slice,
+        step_fn.solver.transport_model,
+        step_fn.solver.source_models,
+        step_fn.solver.pedestal_model,
+        step_fn.solver.neoclassical_models,
+    )
+    self.assertEqual(model1, model2)
+    self.assertEqual(hash(model1), hash(model2))
+
+
+if __name__ == "__main__":
+  absltest.main()

torax/_src/orchestration/step_function.py

@@ -43,6 +43,7 @@
 from torax._src.sources import source_profiles as source_profiles_lib
 from torax._src.time_step_calculator import time_step_calculator as ts
 from torax._src.transport_model import transport_model as transport_model_lib
+import typing_extensions
 
 # pylint: disable=invalid-name
 
@@ -623,6 +624,25 @@ def body_fun(
         geo_t_plus_dt,
     )
 
+  def __hash__(self) -> int:
+    return hash((
+        self.solver,
+        self.mhd_models,
+        self.time_step_calculator,
+        self._dynamic_runtime_params_slice_provider,
+        self._geometry_provider,
+    ))
+
+  def __eq__(self, other: typing_extensions.Self) -> bool:
+    return (
+        self.solver == other.solver
+        and self.mhd_models == other.mhd_models
+        and self.time_step_calculator == other.time_step_calculator
+        and self._dynamic_runtime_params_slice_provider
+        == other._dynamic_runtime_params_slice_provider
+        and self._geometry_provider == other._geometry_provider
+    )
+
 
 @functools.partial(
     jax_utils.jit,

torax/_src/time_step_calculator/chi_time_step_calculator.py

@@ -69,3 +69,9 @@ def next_dt(
     )
 
     return dt
+
+  def __hash__(self) -> int:
+    return hash(self.__class__.__name__)
+
+  def __eq__(self, other) -> bool:
+    return isinstance(other, ChiTimeStepCalculator)

torax/_src/time_step_calculator/fixed_time_step_calculator.py

@@ -56,3 +56,9 @@ def next_dt(
     dt = jnp.array(dynamic_runtime_params_slice.numerics.fixed_dt)
 
     return dt
+
+  def __hash__(self) -> int:
+    return hash(self.__class__.__name__)
+
+  def __eq__(self, other) -> bool:
+    return isinstance(other, FixedTimeStepCalculator)

torax/_src/time_step_calculator/tests/time_step_calculator_test.py

@@ -0,0 +1,43 @@
+# Copyright 2024 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from absl.testing import absltest
+from absl.testing import parameterized
+from torax._src.time_step_calculator import chi_time_step_calculator
+from torax._src.time_step_calculator import fixed_time_step_calculator
+
+
+class TimeStepCalculatorTest(parameterized.TestCase):
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name='fixed',
+          calculator_type=fixed_time_step_calculator.FixedTimeStepCalculator,
+      ),
+      dict(
+          testcase_name='chi',
+          calculator_type=chi_time_step_calculator.ChiTimeStepCalculator,
+      ),
+  )
+  def test_different_time_step_calculators_have_same_hash_and_equals(
+      self,
+      calculator_type,
+  ):
+    calculator1 = calculator_type()
+    calculator2 = calculator_type()
+    self.assertEqual(calculator1, calculator2)
+    self.assertEqual(hash(calculator1), hash(calculator2))
+
+
+if __name__ == '__main__':
+  absltest.main()

torax/_src/time_step_calculator/time_step_calculator.py

@@ -68,3 +68,20 @@ def next_dt(
       core_profiles: Core plasma profiles in the tokamak.
       core_transport: Transport coefficients.
     """
+
+  @abc.abstractmethod
+  def __hash__(self) -> int:
+    """Returns a hash of the time step calculator.
+
+    Should be implemented to support jax.jit caching.
+    """
+
+  @abc.abstractmethod
+  def __eq__(self, other) -> bool:
+    """Returns whether the time step calculator is equal to the other.
+
+    Should be implemented to support jax.jit caching.
+
+    Args:
+      other: The object to compare to.
+    """