Bump up JAX version used in tests

Signed-off-by: Kamil Tokarski <[email protected]>
NVIDIA · Apr 12, 2024 · 8a37bfa · 8a37bfa
1 parent 3b1b5f8
commit 8a37bfa
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/qa/setup_packages.py b/qa/setup_packages.py
@@ -565,9 +565,10 @@ def get_pyvers_name(self, url, cuda_version):
                 PckgVer(
                     "0.4.13", python_min_ver="3.8", python_max_ver="3.8", dependencies=["jaxlib"]
                 ),
-                # dax.fn.jax_function requires 0.4.16 which is the first one supporting
+                # dax.fn.jax_function requires at least 0.4.16 which is the first one supporting
                 # `__dlpack__` method, while 0.4.13 is the last one supported with Python3.8
-                PckgVer("0.4.16", python_min_ver="3.9", dependencies=["jaxlib"]),
+                # Use the most recent one avialable due to issues with address sanitizer run
+                PckgVer("0.4.26", python_min_ver="3.9", dependencies=["jaxlib"]),
             ]
         },
         # name used during installation