Fix jax.config import

tum-pbs · Feb 29, 2024 · 92c162d · 92c162d
1 parent 4cdb4bb
commit 92c162d
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/phiml/backend/jax/_jax_backend.py b/phiml/backend/jax/_jax_backend.py
@@ -20,8 +20,7 @@
 from .._dtype import DType, to_numpy_dtype, from_numpy_dtype
 from .._backend import Backend, ComputeDevice, combined_dim, ML_LOGGER, TensorType, map_structure
 
-from jax.config import config
-config.update("jax_enable_x64", True)
+jax.config.update("jax_enable_x64", True)
 
 
 class JaxBackend(Backend):