From e32ad181c87b11e06fff5b7b56036f3e080d82ee Mon Sep 17 00:00:00 2001 From: nstarman Date: Thu, 18 Jan 2024 16:59:54 -0500 Subject: [PATCH] Fix type hint to mockstream vmap arg Signed-off-by: nstarman --- src/galax/dynamics/mockstream/_mockstream_generator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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(