Fix cache misses on second simulation step due to weakly type Phi_b_dot.

This affects cases with calcphibdot=True.

PiperOrigin-RevId: 834316997

Author: Nush395

torax/_src/geometry/geometry_provider.py

@@ -232,9 +232,11 @@ def _get_geometry_base(
         continue
       if attr.name == 'Phi_b_dot':
         if self.calcphibdot:
-          kwargs[attr.name] = _Phi_b_grad(self.Phi_face, t)
+          kwargs[attr.name] = jnp.asarray(
+              _Phi_b_grad(self.Phi_face, t), dtype=jax_utils.get_dtype()
+          )
         else:
-          kwargs[attr.name] = 0.0
+          kwargs[attr.name] = jnp.zeros((), dtype=jax_utils.get_dtype())
         continue
       provider_attr = getattr(self, attr.name)
       if isinstance(

torax/_src/geometry/standard_geometry.py

@@ -32,6 +32,7 @@
 from torax._src import array_typing
 from torax._src import constants
 from torax._src import interpolated_param
+from torax._src import jax_utils
 from torax._src.geometry import geometry
 from torax._src.geometry import geometry_loader
 from torax._src.geometry import geometry_provider
@@ -1385,7 +1386,7 @@ def build_standard_geometry(
       # always initialize Phibdot as zero. It will be replaced once both geo_t
       # and geo_t_plus_dt are provided, and set to be the same for geo_t and
       # geo_t_plus_dt for each given time interval.
-      Phi_b_dot=np.asarray(0.0),
+      Phi_b_dot=np.zeros((), dtype=jax_utils.get_int_dtype()),
       _z_magnetic_axis=intermediate.z_magnetic_axis,
       diverted=intermediate.diverted,
       connection_length_target=intermediate.connection_length_target,