fix: removed the jax.config import from the ivy_tests init as it's de…

…precated in the recent release (ivy-llc#28444)
fnhirwa · Feb 28, 2024 · 8367486 · 8367486
1 parent 61e3532
commit 8367486
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 16 deletions.
diff --git a/ivy/utils/assertions.py b/ivy/utils/assertions.py
@@ -344,17 +344,14 @@ def check_dev_correct_formatting(device):
 
 
 def _check_jax_x64_flag(dtype):
-    if ivy.backend == "jax":
-        import jax
-
-        if not jax.config.x64_enabled:
-            ivy.utils.assertions.check_elem_in_list(
-                dtype,
-                ["float64", "int64", "uint64", "complex128"],
-                inverse=True,
-                message=(
-                    f"{dtype} output not supported while jax_enable_x64"
-                    " is set to False, please import jax and enable the flag using "
-                    "jax.config.update('jax_enable_x64', True)"
-                ),
-            )
+    if ivy.backend == "jax" and not ivy.functional.backends.jax.jax.config.x64_enabled:
+        ivy.utils.assertions.check_elem_in_list(
+            dtype,
+            ["float64", "int64", "uint64", "complex128"],
+            inverse=True,
+            message=(
+                f"{dtype} output not supported while jax_enable_x64"
+                " is set to False, please import jax and enable the flag using "
+                "jax.config.update('jax_enable_x64', True)"
+            ),
+        )
diff --git a/ivy_tests/__init__.py b/ivy_tests/__init__.py
@@ -1,6 +1,6 @@
 try:
-    from jax.config import config
+    import jax
 
-    config.update("jax_enable_x64", True)
+    jax.config.update("jax_enable_x64", True)
 except (ImportError, RuntimeError):
     pass