Mark n_corrector_steps as static for allowing for backwards mode differentiation of predictor corrector.

jax.lax.fori_loop requires a static upper and lower to be backwards mode differentiable.

PiperOrigin-RevId: 794067640

Author: Nush395

torax/_src/solver/pydantic_model.py

@@ -55,7 +55,9 @@ class BaseSolver(torax_pydantic.BaseModelFrozen, abc.ABC):
       torax_pydantic.UnitInterval, torax_pydantic.JAX_STATIC
   ] = 1.0
   use_predictor_corrector: Annotated[bool, torax_pydantic.JAX_STATIC] = False
-  n_corrector_steps: pydantic.PositiveInt = 10
+  n_corrector_steps: Annotated[
+      pydantic.PositiveInt, torax_pydantic.JAX_STATIC
+  ] = 10
   convection_dirichlet_mode: Annotated[
       Literal['ghost', 'direct', 'semi-implicit'], torax_pydantic.JAX_STATIC
   ] = 'ghost'

torax/_src/solver/runtime_params.py

@@ -36,4 +36,4 @@ class DynamicRuntimeParams:
 
   chi_pereverzev: float
   D_pereverzev: float  # pylint: disable=invalid-name
-  n_corrector_steps: int
+  n_corrector_steps: int = dataclasses.field(metadata={'static': True})