Improve compatibility with jax (#351)

Author: p-zubieta

pysages/typing.py

@@ -8,15 +8,16 @@
 from importlib import import_module
 
 import jax
-import jaxlib.xla_extension as xe
 
 from pysages._compat import _jax_version_tuple, _plum_version_tuple
 
 # Compatibility for jax >=0.4.1
 
 # https://github.com/google/jax/releases/tag/jax-v0.4.1
 if _jax_version_tuple < (0, 4, 1):
+    xe = import_module("jaxlib.xla_extension")
     JaxArray = xe.DeviceArray
+    del xe
 else:
     JaxArray = jax.Array
 
@@ -47,7 +48,6 @@
 
 # Remove namespace noise
 del jax
-del xe
 del import_module
 del _jax_version_tuple
 del _plum_version_tuple