diff --git a/src/galax/dynamics/mockstream/_mockstream_generator.py b/src/galax/dynamics/mockstream/_mockstream_generator.py
index a1cf5afc..08518ca4 100644
--- a/src/galax/dynamics/mockstream/_mockstream_generator.py
+++ b/src/galax/dynamics/mockstream/_mockstream_generator.py
@@ -104,7 +104,7 @@ def _run_vmap(  # TODO: output shape depends on the input shape
         # TODO: make this a separated method
         @jax.jit  # type: ignore[misc]
         def single_particle_integrate(
-            i: int, qp0_lead_i: Vec6, qp0_trail_i: Vec6
+            i: IntScalar, qp0_lead_i: Vec6, qp0_trail_i: Vec6
         ) -> tuple[Vec6, Vec6]:
             t_i = ts[i]
             qp_lead = self.potential.integrate_orbit(