Move usages of functools.partial(jax.jit, ...) to use decorator factory pattern.

PiperOrigin-RevId: 836233292

Author: Nush395

docs/jax_classes.rst

@@ -87,7 +87,7 @@ Designing JAX-compatible classes
     import functools
     import jax
 
-    @functools.partial(jax.jit, static_argnums=0)
+    @jax.jit(static_argnums=0)
     def f(x):
       return x.num_wheels
 

torax/_src/core_profiles/updaters.py

@@ -28,7 +28,6 @@
   the next timestep and returns updates to State.
 """
 import dataclasses
-import functools
 
 import jax
 from jax import numpy as jnp
@@ -135,7 +134,7 @@ def get_prescribed_core_profile_values(
   }
 
 
-@functools.partial(jax.jit, static_argnames=['evolving_names'])
+@jax.jit(static_argnames='evolving_names')
 def update_core_profiles_during_step(
     x_new: tuple[cell_variable.CellVariable, ...],
     runtime_params: runtime_params_lib.RuntimeParams,

torax/_src/orchestration/jit_run_loop.py

@@ -13,8 +13,6 @@
 # limitations under the License.
 
 """JITted run_loop for iterating over the simulation step function."""
-import functools
-
 import chex
 import jax
 import jax.numpy as jnp
@@ -27,7 +25,7 @@
 from torax._src.output_tools import post_processing
 
 
-@functools.partial(jax.jit, static_argnames=['max_steps'])
+@jax.jit(static_argnames='max_steps')
 def run_loop_jit(
     step_fn: step_function.SimulationStepFn,
     max_steps: int,

torax/_src/solver/tests/jax_root_finding_test.py

@@ -52,7 +52,7 @@ def test_root_newton_raphson_basic(self):
     x_init = np.array([0.0, 0.0], dtype=dtype)
     sol_np = optimize.root(f_closed, [0, 0], tol=tol)
 
-    @functools.partial(jax.jit, static_argnames=['tol', 'maxiter'])
+    @jax.jit(static_argnames=['tol', 'maxiter'])
     def root_jax(x, tol, maxiter):
       return jax_root_finding.root_newton_raphson(
           f_closed, x, tol=tol, maxiter=maxiter

torax/_src/sources/ion_cyclotron_source.py

@@ -266,7 +266,7 @@ def __eq__(self, other: typing_extensions.Self) -> bool:
     return isinstance(other, ToricNNWrapper)
 
 
-@functools.partial(jax.jit, static_argnames='toric_nn')
+@jax.jit(static_argnames='toric_nn')
 def _toric_nn_predict(
     toric_nn: ToricNNWrapper,
     inputs: ToricNNInputs,

torax/_src/tests/jax_utils_test.py

@@ -11,8 +11,6 @@
 # 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 functools
 import os
 from unittest import mock
 
@@ -134,7 +132,7 @@ def f(x: jax.Array):
   @parameterized.parameters(['while_loop', 'pure_callback'])
   def test_non_inlined_function(self, implementation):
 
-    @functools.partial(jax.jit, static_argnames=['z'])
+    @jax.jit(static_argnames='z')
     def f(x, z, y=2.0):
       if z == 'left':
         return x['temp1'] * y + jnp.sin(x['temp2'])

torax/_src/transport_model/transport_coefficients_builder.py

@@ -14,7 +14,6 @@
 
 """Code to build the combined transport coefficients for a simulation."""
 import dataclasses
-import functools
 
 import jax
 from torax._src import state
@@ -25,7 +24,7 @@
 from torax._src.transport_model import transport_model as transport_model_lib
 
 
-@functools.partial(jax.jit, static_argnums=(0, 1, 2))
+@jax.jit(static_argnums=(0, 1, 2))
 def calculate_total_transport_coeffs(
     pedestal_model: pedestal_model_lib.PedestalModel,
     transport_model: transport_model_lib.TransportModel,