Modify JIT of transport model to be defined on hash of TransportModel so that subsequent simulations hit cache.

Also:
- remove jitted_transport_model from step function as this will be cache missed in subsequent simulations.
- remove inner JIT from qlknn_transport_model as this is now covered by the JIT on __call__.

Otherwise this misses a lookup as the function is defined on an object by id.

PiperOrigin-RevId: 739945207

Author: Nush395

torax/orchestration/step_function.py

@@ -70,9 +70,6 @@ def __init__(
     self._time_step_calculator = time_step_calculator
     self._transport_model = transport_model
     self._pedestal_model = pedestal_model
-    self._jitted_transport_model = jax_utils.jit(
-        transport_model.__call__,
-    )
 
   @property
   def pedestal_model(self) -> pedestal_model_lib.PedestalModel:
@@ -239,7 +236,7 @@ def init_time_step_calculator(
     pedestal_model_output = self._pedestal_model(
         dynamic_runtime_params_slice_t, geo_t, input_state.core_profiles
     )
-    transport_coeffs = self._jitted_transport_model(
+    transport_coeffs = self._transport_model(
         dynamic_runtime_params_slice_t,
         geo_t,
         input_state.core_profiles,

torax/transport_model/qlknn_transport_model.py

@@ -212,12 +212,6 @@ def _call_implementation(
     )
     return self._combined(runtime_config_inputs, geo, core_profiles)
 
-  # Wrap in JIT here in order to cache the tracing/compilation of this function.
-  # We mark self as static because it is a singleton. Other args are pytrees.
-  # There's no global coordination of calls to transport model so it is called
-  # 2-4X with the same args. Caching prevents construction of multiple copies of
-  # identical expressions saving ~30% in compile time.
-  @functools.partial(jax.jit, static_argnames=['self'])
   def _combined(
       self,
       runtime_config_inputs: QLKNNRuntimeConfigInputs,

torax/transport_model/tests/qlknn_transport_model_test.py

@@ -76,15 +76,15 @@ def test_qlknn_transport_model_cache_works(self):
     qlknn(
         dynamic_runtime_params_slice, geo, core_profiles, pedestal_model_outputs
     )
-    cache_size = qlknn._combined._cache_size()  # pylint: disable=protected-access
+    cache_size = qlknn.__call__._cache_size()  # pytype: disable=attribute-error
     self.assertGreaterEqual(cache_size, 1)
 
     # Executing again should lead to the same cache entry being used.
     qlknn(
         dynamic_runtime_params_slice, geo, core_profiles, pedestal_model_outputs
     )
     self.assertEqual(
-        qlknn._combined._cache_size(),  # pylint: disable=protected-access
+        qlknn.__call__._cache_size(),  # pytype: disable=attribute-error
         cache_size,
     )
 

torax/transport_model/transport_model.py

@@ -20,10 +20,12 @@
 
 import abc
 import dataclasses
+import functools
 
 import jax
 from jax import numpy as jnp
 from torax import constants
+from torax import jax_utils
 from torax import state
 from torax.config import runtime_params_slice
 from torax.geometry import geometry
@@ -53,6 +55,7 @@ def __setattr__(self, attr, value):
       raise AttributeError("TransportModels are immutable.")
     return super().__setattr__(attr, value)
 
+  @functools.partial(jax_utils.jit, static_argnums=(0,))
   def __call__(
       self,
       dynamic_runtime_params_slice: runtime_params_slice.DynamicRuntimeParamsSlice,