changed handling of jax dependency

Jashcraf · Aug 16, 2024 · 0d6619b · 0d6619b
1 parent dbc4959
commit 0d6619b
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/tests/test_katsu_math.py b/tests/test_katsu_math.py
@@ -1,4 +1,3 @@
-import jax.numpy as jnp
 import pytest
 from katsu.katsu_math import (
     BackendShim,
@@ -18,7 +17,7 @@
     cupy_installed = False
 
 try:
-    import jax.numpy
+    import jax.numpy as jnp
     jax_installed = True
 
 except Exception: