jax wrapper

Author: Jammy2211

autoconf/jax_wrapper.py

@@ -45,6 +45,7 @@
 
     import jax
     from jax import numpy
+    from jax import numpy as np
 
     jax_enable_x64 = jax.config.read("jax_enable_x64")
 
@@ -89,6 +90,7 @@ def grad(function, *args, **kwargs):
     """)
 
     import numpy  # noqa
+    import numpy as np
 
     def jit(function, *_, **__):
         return function