Fix type hints

Author: amacati

tests/unit/test_functional.py

@@ -12,6 +12,8 @@
 from crazyflow.sim import Sim
 
 if TYPE_CHECKING:
+    from jax import Array
+
     from crazyflow.sim.data import SimData
 
 
@@ -89,7 +91,7 @@ def test_functional_attitude_control(attitude_freq: int):
     # Check if we can apply inside of jax jit which does not permit device tracking. See
     # https://github.com/jax-ml/jax/issues/26000 for more context.
     @jax.jit
-    def apply_control(data: SimData, cmd: jnp.ndarray) -> SimData:
+    def apply_control(data: SimData, cmd: Array) -> SimData:
         return F.attitude_control(data, cmd)
 
     jax.block_until_ready(apply_control(data, cmd))
@@ -103,7 +105,7 @@ def test_functional_attitude_control_device(device: str):
     cmd = np.random.rand(sim.n_worlds, sim.n_drones, 4)
     data = F.attitude_control(data, cmd)
     controls = data.controls.attitude
-    assert isinstance(controls.staged_cmd, jnp.ndarray), "Buffers must remain JAX arrays"
+    assert isinstance(controls.staged_cmd, Array), "Buffers must remain JAX arrays"
     assert jnp.all(controls.staged_cmd == cmd), "Buffers must match command"