Fix type hint to mockstream vmap arg (GalacticDynamics#65)

Signed-off-by: nstarman <nstarman@users.noreply.github.com>
nstarman · Jan 18, 2024 · e309553 · e309553
1 parent fb2a422
commit e309553
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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(