Add experimental compile to optimizer solver.

Speed test on implicit_short_optimizer config
Before:
First sim: 13.72s
Second sim: 0.23s

After w. experimental compile:
First sim: 13.18s
Second sim: 0.22s

After w/o experimental compile:
First sim: 13.71s
Second sim: 0.21s

PiperOrigin-RevId: 774765146

Author: Nush395

torax/_src/fvm/optimizer_solve_block.py

@@ -15,11 +15,12 @@
 
 See function docstring for details.
 """
-
+import functools
 from typing import TypeAlias
 
 import jax
 from torax._src import state
+from torax._src import xnp
 from torax._src.config import runtime_params_slice
 from torax._src.core_profiles import convertors
 from torax._src.fvm import block_1d_coeffs
@@ -39,6 +40,19 @@
 AuxiliaryOutput: TypeAlias = block_1d_coeffs.AuxiliaryOutput
 
 
+@functools.partial(
+    xnp.jit,
+    static_argnames=[
+        'static_runtime_params_slice',
+        'transport_model',
+        'source_models',
+        'neoclassical_models',
+        'pedestal_model',
+        'coeffs_callback',
+        'evolving_names',
+        'initial_guess_mode',
+    ],
+)
 def optimizer_solve_block(
     dt: jax.Array,
     static_runtime_params_slice: runtime_params_slice.StaticRuntimeParamsSlice,

torax/tests/sim_experimental_compile_test.py

@@ -41,6 +41,11 @@ class SimExperimentalCompileTest(sim_test_case.SimTestCase):
           'test_iterhybrid_predictor_corrector',
           'test_iterhybrid_predictor_corrector.py',
       ),
+      # Using optimizer solver.
+      (
+          'test_implicit_short_optimizer',
+          'test_implicit_short_optimizer.py',
+      ),
   )
   def test_run_simulation_experimental_compile(
       self,