{ai}[gfbf/2025a] jax v0.6.2 by pavelToman · Pull Request #23526 · easybuilders/easybuild-easyconfigs

pavelToman · 2025-07-29T18:44:10Z

(created using eb --new-pr)
resolves vscentrum/vsc-software-stack#578
requires:

Competing PR jax-0.7.0:

{ai}[gfbf/2025a] jax v0.7.0 #23530

To pass tests it needs more memory - at least 50gb memory with 8 cpus.

…rsion.patch

github-actions · 2025-07-29T18:44:39Z

Updated software `jax-0.6.2-gfbf-2025a.eb`

Diff against jax-0.4.25-gfbf-2023a.eb

easybuild/easyconfigs/j/jax/jax-0.4.25-gfbf-2023a.eb

diff --git a/easybuild/easyconfigs/j/jax/jax-0.4.25-gfbf-2023a.eb b/easybuild/easyconfigs/j/jax/jax-0.6.2-gfbf-2025a.eb
index e61bc4719d..6b91637431 100644
--- a/easybuild/easyconfigs/j/jax/jax-0.4.25-gfbf-2023a.eb
+++ b/easybuild/easyconfigs/j/jax/jax-0.6.2-gfbf-2025a.eb
@@ -6,99 +6,86 @@
 easyblock = 'PythonBundle'
 
 name = 'jax'
-version = '0.4.25'
+version = '0.6.2'
 
 homepage = 'https://jax.readthedocs.io/'
 description = """Composable transformations of Python+NumPy programs:
 differentiate, vectorize, JIT to GPU/TPU, and more"""
 
-toolchain = {'name': 'gfbf', 'version': '2023a'}
+toolchain = {'name': 'gfbf', 'version': '2025a'}
 
 builddependencies = [
-    ('Bazel', '6.3.1'),
-    ('pytest-xdist', '3.3.1'),
-    ('git', '2.41.0', '-nodocs'),  # bazel uses git to fetch repositories
-    ('matplotlib', '3.7.2'),  # required for tests/lobpcg_test.py
-    ('poetry', '1.5.1'),
-    ('pybind11', '2.11.1'),
+    ('Bazel', '7.4.1', '-Java-21'),
 ]
 
 dependencies = [
-    ('Python', '3.11.3'),
-    ('SciPy-bundle', '2023.07'),
-    ('absl-py', '2.1.0'),
-    ('flatbuffers-python', '23.5.26'),
-    ('ml_dtypes', '0.3.2'),
-    ('zlib', '1.2.13'),
+    ('Python', '3.13.1'),
+    ('SciPy-bundle', '2025.06'),
+    ('absl-py', '2.3.1'),
+    ('flatbuffers-python', '25.2.10'),
+    ('ml_dtypes', '0.5.1'),
+    ('hypothesis', '6.133.2'),
+    ('zlib', '1.3.1'),
 ]
 
-# downloading xla and other tarballs to avoid that Bazel downloads it during the build
-local_extract_cmd = 'mkdir -p %(builddir)s/archives && cp %s %(builddir)s/archives'
+# downloading xla  tarball to avoid that Bazel downloads it during the build
 # note: following commits *must* be the exact same onces used upstream
 # XLA_COMMIT from jax-jaxlib: third_party/xla/workspace.bzl
-local_xla_commit = '4ccfe33c71665ddcbca5b127fefe8baa3ed632d4'
-# TFRT_COMMIT from xla: third_party/tsl/third_party/tf_runtime/workspace.bzl
-local_tfrt_commit = '0aeefb1660d7e37964b2bb71b1f518096bda9a25'
+_xla_commit = '3d5ece64321630dade7ff733ae1353fc3c83d9cc'
 
 # Use sources downloaded by EasyBuild
 _jaxlib_buildopts = '--bazel_options="--distdir=%(builddir)s/archives" '
-# Use dependencies from EasyBuild
-_jaxlib_buildopts += '--bazel_options="--action_env=TF_SYSTEM_LIBS=pybind11" '
-_jaxlib_buildopts += '--bazel_options="--action_env=CPATH=$EBROOTPYBIND11/include" '
-# Avoid warning (treated as error) in upb/table.c
-_jaxlib_buildopts += '--bazel_options="--copt=-Wno-maybe-uninitialized" '
+# create wheels for jaxlib and install it
+_jaxlib_buildopts += '--wheels=jaxlib '
+# use GCC instead of default Clang
+_jaxlib_buildopts += '--use_clang=false --gcc_path=$CC '
+# fix jaxlib version
+_jaxlib_buildopts += '--bazel_options="--action_env=JAXLIB_RELEASE" '
 
 components = [
     ('jaxlib', version, {
         'sources': [
             {
                 'source_urls': ['https://github.com/google/jax/archive/'],
-                'filename': '%(name)s-v%(version)s.tar.gz',
+                'filename': 'jax-v%(version)s.tar.gz',
             },
             {
                 'source_urls': ['https://github.com/openxla/xla/archive'],
-                'download_filename': '%s.tar.gz' % local_xla_commit,
-                'filename': 'xla-%s.tar.gz' % local_xla_commit[:8],
-                'extract_cmd': local_extract_cmd,
-            },
-            {
-                'source_urls': ['https://github.com/tensorflow/runtime/archive'],
-                'download_filename': '%s.tar.gz' % local_tfrt_commit,
-                'filename': 'tf_runtime-%s.tar.gz' % local_tfrt_commit[:8],
-                'extract_cmd': local_extract_cmd,
+                'download_filename': f'{_xla_commit}.tar.gz',
+                'filename': f'xla-{_xla_commit[:8]}.tar.gz',
+                'extract_cmd': 'mkdir -p %(builddir)s/archives && cp %s %(builddir)s/archives',
             },
         ],
-        'patches': ['jax-0.4.25_fix-pybind11-systemlib.patch'],
         'checksums': [
-            {'jaxlib-v0.4.25.tar.gz':
-             'fc1197c401924942eb14185a61688d0c476e3e81ff71f9dc95e620b57c06eec8'},
-            {'xla-4ccfe33c.tar.gz':
-             '8a59b9af7d0850059d7043f7043c780066d61538f3af536e8a10d3d717f35089'},
-            {'tf_runtime-0aeefb16.tar.gz':
-             'a3df827d7896774cb1d80bf4e1c79ab05c268f29bd4d3db1fb5a4b9c2079d8e3'},
-            {'jax-0.4.25_fix-pybind11-systemlib.patch':
-             'daad5b726d1a138431b05eb60ecf4c89c7b5148eb939721800bdf43d804ca033'},
+            {'jax-v0.6.2.tar.gz':
+             'd46cb98795f2c1ccdf2b081e02d9d74b659063679a80beb001ad17d482a60e17'},
+            {'xla-3d5ece64.tar.gz':
+             'fbd20cf83bad78f66977fa7ff67a12e52964abae0b107ddd5486a0355643ec8a'},
         ],
-        'start_dir': 'jax-jaxlib-v%(version)s',
-        'buildopts': _jaxlib_buildopts
+        'start_dir': f'jax-jax-v{version}',
+        # fix jaxlib version - removes .dev suffix
+        'prebuildopts': 'export JAXLIB_RELEASE=%(version)s && ',
+        'buildopts': _jaxlib_buildopts,
     }),
 ]
+# BackendsTest is failing on systems with GPU detected when cpu version is built
+_failing_backend_test = 'tests/api_test.py::BackendsTest::test_no_backend_warning_on_cpu_if_platform_specified'
+# Failing after fix jaxlib version by JAXLIB_RELEASE
+_failing_version_tests = 'tests/version_test.py'
+_runtest_cmd = f"pytest -n %(parallel)s tests --deselect={_failing_backend_test} --deselect={_failing_version_tests}"
 
 exts_list = [
     (name, version, {
-        'sources': [
-            {
-                'source_urls': ['https://github.com/google/jax/archive/'],
-                'filename': '%(name)s-v%(version)s.tar.gz',
-            },
+        "patches": ["jax-0.6.2_version.patch"],
+        "testinstall": True,
+        "runtest": _runtest_cmd,
+        "sources": [
+            {"source_urls": ["https://github.com/google/jax/archive/"], "filename": "%(name)s-v%(version)s.tar.gz"}
         ],
-        'patches': ['jax-0.4.25_fix_env_test_no_log_spam.patch'],
-        'checksums': [
-            {'jax-v0.4.25.tar.gz': '8b30af49688c0c13b82c6f5ce992727c00b5fc6d04a4c6962012f4246fa664eb'},
-            {'jax-0.4.25_fix_env_test_no_log_spam.patch':
-             'a18b5f147569d9ad41025124333a0f04fd0d0e0f9e4309658d7f6b9b838e2e2a'},
+        "checksums": [
+            {"jax-v0.6.2.tar.gz": "d46cb98795f2c1ccdf2b081e02d9d74b659063679a80beb001ad17d482a60e17"},
+            {"jax-0.6.2_version.patch": "e15615fd9f4e1698f7c5fd384f146d7b2dbfde3d4657b69bd2d044d75c9fb1d4"},
         ],
-        'runtest': "pytest -n %(parallel)s tests",
     }),
 ]

Diff against jax-0.4.25-gfbf-2023a-CUDA-12.1.1.eb

easybuild/easyconfigs/j/jax/jax-0.4.25-gfbf-2023a-CUDA-12.1.1.eb

diff --git a/easybuild/easyconfigs/j/jax/jax-0.4.25-gfbf-2023a-CUDA-12.1.1.eb b/easybuild/easyconfigs/j/jax/jax-0.6.2-gfbf-2025a.eb
index e0c599039d..6b91637431 100644
--- a/easybuild/easyconfigs/j/jax/jax-0.4.25-gfbf-2023a-CUDA-12.1.1.eb
+++ b/easybuild/easyconfigs/j/jax/jax-0.6.2-gfbf-2025a.eb
@@ -6,125 +6,86 @@
 easyblock = 'PythonBundle'
 
 name = 'jax'
-version = '0.4.25'
-versionsuffix = '-CUDA-%(cudaver)s'
+version = '0.6.2'
 
 homepage = 'https://jax.readthedocs.io/'
 description = """Composable transformations of Python+NumPy programs:
 differentiate, vectorize, JIT to GPU/TPU, and more"""
 
-toolchain = {'name': 'gfbf', 'version': '2023a'}
-cuda_compute_capabilities = ["5.0", "6.0", "6.1", "7.0", "7.5", "8.0", "8.6", "9.0"]
+toolchain = {'name': 'gfbf', 'version': '2025a'}
 
 builddependencies = [
-    ('Bazel', '6.3.1'),
-    ('pytest-xdist', '3.3.1'),
-    ('git', '2.41.0', '-nodocs'),  # bazel uses git to fetch repositories
-    ('matplotlib', '3.7.2'),  # required for tests/lobpcg_test.py
-    ('poetry', '1.5.1'),
-    ('pybind11', '2.11.1'),
+    ('Bazel', '7.4.1', '-Java-21'),
 ]
 
 dependencies = [
-    ('CUDA', '12.1.1', '', SYSTEM),
-    ('cuDNN', '8.9.2.26', versionsuffix, SYSTEM),
-    ('NCCL', '2.18.3', versionsuffix),
-    ('Python', '3.11.3'),
-    ('SciPy-bundle', '2023.07'),
-    ('absl-py', '2.1.0'),
-    ('flatbuffers-python', '23.5.26'),
-    ('ml_dtypes', '0.3.2'),
-    ('zlib', '1.2.13'),
+    ('Python', '3.13.1'),
+    ('SciPy-bundle', '2025.06'),
+    ('absl-py', '2.3.1'),
+    ('flatbuffers-python', '25.2.10'),
+    ('ml_dtypes', '0.5.1'),
+    ('hypothesis', '6.133.2'),
+    ('zlib', '1.3.1'),
 ]
 
-# downloading xla and other tarballs to avoid that Bazel downloads it during the build
-local_extract_cmd = 'mkdir -p %(builddir)s/archives && cp %s %(builddir)s/archives'
+# downloading xla  tarball to avoid that Bazel downloads it during the build
 # note: following commits *must* be the exact same onces used upstream
 # XLA_COMMIT from jax-jaxlib: third_party/xla/workspace.bzl
-local_xla_commit = '4ccfe33c71665ddcbca5b127fefe8baa3ed632d4'
-# TFRT_COMMIT from xla: third_party/tsl/third_party/tf_runtime/workspace.bzl
-local_tfrt_commit = '0aeefb1660d7e37964b2bb71b1f518096bda9a25'
+_xla_commit = '3d5ece64321630dade7ff733ae1353fc3c83d9cc'
 
 # Use sources downloaded by EasyBuild
 _jaxlib_buildopts = '--bazel_options="--distdir=%(builddir)s/archives" '
-# Use dependencies from EasyBuild
-_jaxlib_buildopts += '--bazel_options="--action_env=TF_SYSTEM_LIBS=pybind11" '
-_jaxlib_buildopts += '--bazel_options="--action_env=CPATH=$EBROOTPYBIND11/include" '
-# Avoid warning (treated as error) in upb/table.c
-_jaxlib_buildopts += '--bazel_options="--copt=-Wno-maybe-uninitialized" '
-
-# Some tests require an isolated run:
-local_isolated_tests = [
-    'tests/host_callback_test.py::HostCallbackTapTest::test_tap_scan_custom_jvp',
-    'tests/host_callback_test.py::HostCallbackTapTest::test_tap_transforms_doc',
-    'tests/lax_scipy_special_functions_test.py::LaxScipySpcialFunctionsTest' +
-    '::testScipySpecialFun_gammainc_s_2x1x4_float32_float32',
-]
-# deliberately not testing in parallel, as that results in (additional) failing tests;
-# use XLA_PYTHON_CLIENT_ALLOCATOR=platform to allocate and deallocate GPU memory during testing,
-# see https://github.com/google/jax/issues/7323 and
-# https://github.com/google/jax/blob/main/docs/gpu_memory_allocation.rst;
-# use CUDA_VISIBLE_DEVICES=0 to avoid failing tests on systems with multiple GPUs;
-# use NVIDIA_TF32_OVERRIDE=0 to avoid loosing numerical precision by disabling TF32 Tensor Cores;
-local_test_exports = [
-    "NVIDIA_TF32_OVERRIDE=0",
-    "CUDA_VISIBLE_DEVICES=0",
-    "XLA_PYTHON_CLIENT_ALLOCATOR=platform",
-    "JAX_ENABLE_X64=true",
-]
-local_test = ''.join(['export %s;' % x for x in local_test_exports])
-# run all tests at once except for local_isolated_tests:
-local_test += "pytest -vv tests %s && " % ' '.join(['--deselect %s' % x for x in local_isolated_tests])
-# run remaining local_isolated_tests separately:
-local_test += ' && '.join(['pytest -vv %s' % x for x in local_isolated_tests])
+# create wheels for jaxlib and install it
+_jaxlib_buildopts += '--wheels=jaxlib '
+# use GCC instead of default Clang
+_jaxlib_buildopts += '--use_clang=false --gcc_path=$CC '
+# fix jaxlib version
+_jaxlib_buildopts += '--bazel_options="--action_env=JAXLIB_RELEASE" '
 
 components = [
     ('jaxlib', version, {
         'sources': [
             {
                 'source_urls': ['https://github.com/google/jax/archive/'],
-                'filename': '%(name)s-v%(version)s.tar.gz',
+                'filename': 'jax-v%(version)s.tar.gz',
             },
             {
                 'source_urls': ['https://github.com/openxla/xla/archive'],
-                'download_filename': '%s.tar.gz' % local_xla_commit,
-                'filename': 'xla-%s.tar.gz' % local_xla_commit[:8],
-                'extract_cmd': local_extract_cmd,
-            },
-            {
-                'source_urls': ['https://github.com/tensorflow/runtime/archive'],
-                'download_filename': '%s.tar.gz' % local_tfrt_commit,
-                'filename': 'tf_runtime-%s.tar.gz' % local_tfrt_commit[:8],
-                'extract_cmd': local_extract_cmd,
+                'download_filename': f'{_xla_commit}.tar.gz',
+                'filename': f'xla-{_xla_commit[:8]}.tar.gz',
+                'extract_cmd': 'mkdir -p %(builddir)s/archives && cp %s %(builddir)s/archives',
             },
         ],
-        'patches': ['jax-0.4.25_fix-pybind11-systemlib.patch'],
         'checksums': [
-            {'jaxlib-v0.4.25.tar.gz':
-             'fc1197c401924942eb14185a61688d0c476e3e81ff71f9dc95e620b57c06eec8'},
-            {'xla-4ccfe33c.tar.gz':
-             '8a59b9af7d0850059d7043f7043c780066d61538f3af536e8a10d3d717f35089'},
-            {'tf_runtime-0aeefb16.tar.gz':
-             'a3df827d7896774cb1d80bf4e1c79ab05c268f29bd4d3db1fb5a4b9c2079d8e3'},
-            {'jax-0.4.25_fix-pybind11-systemlib.patch':
-             'daad5b726d1a138431b05eb60ecf4c89c7b5148eb939721800bdf43d804ca033'},
+            {'jax-v0.6.2.tar.gz':
+             'd46cb98795f2c1ccdf2b081e02d9d74b659063679a80beb001ad17d482a60e17'},
+            {'xla-3d5ece64.tar.gz':
+             'fbd20cf83bad78f66977fa7ff67a12e52964abae0b107ddd5486a0355643ec8a'},
         ],
-        'start_dir': 'jax-jaxlib-v%(version)s',
-        'buildopts': _jaxlib_buildopts
+        'start_dir': f'jax-jax-v{version}',
+        # fix jaxlib version - removes .dev suffix
+        'prebuildopts': 'export JAXLIB_RELEASE=%(version)s && ',
+        'buildopts': _jaxlib_buildopts,
     }),
 ]
+# BackendsTest is failing on systems with GPU detected when cpu version is built
+_failing_backend_test = 'tests/api_test.py::BackendsTest::test_no_backend_warning_on_cpu_if_platform_specified'
+# Failing after fix jaxlib version by JAXLIB_RELEASE
+_failing_version_tests = 'tests/version_test.py'
+_runtest_cmd = f"pytest -n %(parallel)s tests --deselect={_failing_backend_test} --deselect={_failing_version_tests}"
 
 exts_list = [
     (name, version, {
-        'source_tmpl': '%(name)s-v%(version)s.tar.gz',
-        'source_urls': ['https://github.com/google/jax/archive/'],
-        'patches': ['jax-0.4.25_fix_env_test_no_log_spam.patch'],
-        'checksums': [
-            {'jax-v0.4.25.tar.gz': '8b30af49688c0c13b82c6f5ce992727c00b5fc6d04a4c6962012f4246fa664eb'},
-            {'jax-0.4.25_fix_env_test_no_log_spam.patch':
-             'a18b5f147569d9ad41025124333a0f04fd0d0e0f9e4309658d7f6b9b838e2e2a'},
+        "patches": ["jax-0.6.2_version.patch"],
+        "testinstall": True,
+        "runtest": _runtest_cmd,
+        "sources": [
+            {"source_urls": ["https://github.com/google/jax/archive/"], "filename": "%(name)s-v%(version)s.tar.gz"}
+        ],
+        "checksums": [
+            {"jax-v0.6.2.tar.gz": "d46cb98795f2c1ccdf2b081e02d9d74b659063679a80beb001ad17d482a60e17"},
+            {"jax-0.6.2_version.patch": "e15615fd9f4e1698f7c5fd384f146d7b2dbfde3d4657b69bd2d044d75c9fb1d4"},
         ],
-        'runtest': local_test,
     }),
 ]

Diff against jax-0.3.25-foss-2022a.eb

easybuild/easyconfigs/j/jax/jax-0.3.25-foss-2022a.eb

diff --git a/easybuild/easyconfigs/j/jax/jax-0.3.25-foss-2022a.eb b/easybuild/easyconfigs/j/jax/jax-0.6.2-gfbf-2025a.eb
index 27458b716a..6b91637431 100644
--- a/easybuild/easyconfigs/j/jax/jax-0.3.25-foss-2022a.eb
+++ b/easybuild/easyconfigs/j/jax/jax-0.6.2-gfbf-2025a.eb
@@ -1,109 +1,92 @@
 # This file is an EasyBuild reciPY as per https://github.com/easybuilders/easybuild
 # Author: Denis Kristak
 # Updated by: Alex Domingo (Vrije Universiteit Brussel)
+# Updated by: Pavel Tománek (INUITS)
+# Updated by: Thomas Hoffmann (EMBL Heidelberg)
 easyblock = 'PythonBundle'
 
 name = 'jax'
-version = '0.3.25'
+version = '0.6.2'
 
-homepage = 'https://pypi.python.org/pypi/jax'
+homepage = 'https://jax.readthedocs.io/'
 description = """Composable transformations of Python+NumPy programs:
 differentiate, vectorize, JIT to GPU/TPU, and more"""
 
-toolchain = {'name': 'foss', 'version': '2022a'}
+toolchain = {'name': 'gfbf', 'version': '2025a'}
 
 builddependencies = [
-    ('Bazel', '5.1.1'),
-    ('pytest-xdist', '2.5.0'),
-    # git 2.x required to fetch repository 'io_bazel_rules_docker'
-    ('git', '2.36.0', '-nodocs'),
-    ('matplotlib', '3.5.2'),  # required for tests/lobpcg_test.py
+    ('Bazel', '7.4.1', '-Java-21'),
 ]
 
 dependencies = [
-    ('Python', '3.10.4'),
-    ('SciPy-bundle', '2022.05'),
-    ('flatbuffers-python', '2.0'),
+    ('Python', '3.13.1'),
+    ('SciPy-bundle', '2025.06'),
+    ('absl-py', '2.3.1'),
+    ('flatbuffers-python', '25.2.10'),
+    ('ml_dtypes', '0.5.1'),
+    ('hypothesis', '6.133.2'),
+    ('zlib', '1.3.1'),
 ]
 
-# downloading TensorFlow tarball to avoid that Bazel downloads it during the build
-# note: this *must* be the exact same commit as used in WORKSPACE
-local_tf_commit = 'f0fe8d4c04fab1f157854a1aa3c136377901cdef'
-local_tf_dir = 'tensorflow-%s' % local_tf_commit
-local_tf_builddir = '%(builddir)s/' + local_tf_dir
+# downloading xla  tarball to avoid that Bazel downloads it during the build
+# note: following commits *must* be the exact same onces used upstream
+# XLA_COMMIT from jax-jaxlib: third_party/xla/workspace.bzl
+_xla_commit = '3d5ece64321630dade7ff733ae1353fc3c83d9cc'
 
-# replace remote TensorFlow repository with the local one from EB
-local_jax_prebuildopts = "sed -i -f jaxlib_local-tensorflow-repo.sed WORKSPACE && "
-local_jax_prebuildopts += "sed -i 's|EB_TF_REPOPATH|%s|' WORKSPACE && " % local_tf_builddir
-
-default_easyblock = 'PythonPackage'
-default_component_specs = {
-    'sources': [SOURCE_TAR_GZ],
-    'source_urls': [PYPI_SOURCE],
-    'start_dir': '%(name)s-%(version)s',
-}
+# Use sources downloaded by EasyBuild
+_jaxlib_buildopts = '--bazel_options="--distdir=%(builddir)s/archives" '
+# create wheels for jaxlib and install it
+_jaxlib_buildopts += '--wheels=jaxlib '
+# use GCC instead of default Clang
+_jaxlib_buildopts += '--use_clang=false --gcc_path=$CC '
+# fix jaxlib version
+_jaxlib_buildopts += '--bazel_options="--action_env=JAXLIB_RELEASE" '
 
 components = [
-    ('absl-py', '1.3.0', {
-        'options': {'modulename': 'absl'},
-        'checksums': ['463c38a08d2e4cef6c498b76ba5bd4858e4c6ef51da1a5a1f27139a022e20248'],
-    }),
     ('jaxlib', version, {
         'sources': [
-            '%(name)s-v%(version)s.tar.gz',
             {
-                'download_filename': '%s.tar.gz' % local_tf_commit,
-                'filename': 'tensorflow-%s.tar.gz' % local_tf_commit,
-            }
-        ],
-        'source_urls': [
-            'https://github.com/google/jax/archive/',
-            'https://github.com/tensorflow/tensorflow/archive/'
-        ],
-        'patches': [
-            ('jaxlib_local-tensorflow-repo.sed', '.'),
-            ('TensorFlow-2.7.0_cuda-noncanonical-include-paths.patch', '../' + local_tf_dir),
+                'source_urls': ['https://github.com/google/jax/archive/'],
+                'filename': 'jax-v%(version)s.tar.gz',
+            },
+            {
+                'source_urls': ['https://github.com/openxla/xla/archive'],
+                'download_filename': f'{_xla_commit}.tar.gz',
+                'filename': f'xla-{_xla_commit[:8]}.tar.gz',
+                'extract_cmd': 'mkdir -p %(builddir)s/archives && cp %s %(builddir)s/archives',
+            },
         ],
         'checksums': [
-            # jaxlib-v0.3.25.tar.gz
-            '73ebc7868631cd9d520385557bbd7f08762d748a5a6a1bebef0f3b8d7ba748ef',
-            # tensorflow-f0fe8d4c04fab1f157854a1aa3c136377901cdef.tar.gz
-            '9ebba3031e8a81993682e4b9e43891ebb8480b6287e635df8e7efaa45ab5ede7',
-            # jaxlib_local-tensorflow-repo.sed
-            'abb5c3b97f4e317bce9f22ed3eeea3b9715365818d8b50720d937e2d41d5c4e5',
-            # TensorFlow-2.7.0_cuda-noncanonical-include-paths.patch
-            '0a759010c253d49755955cd5f028e75de4a4c447dcc8f5a0d9f47cce6881a9db',
+            {'jax-v0.6.2.tar.gz':
+             'd46cb98795f2c1ccdf2b081e02d9d74b659063679a80beb001ad17d482a60e17'},
+            {'xla-3d5ece64.tar.gz':
+             'fbd20cf83bad78f66977fa7ff67a12e52964abae0b107ddd5486a0355643ec8a'},
         ],
-        'start_dir': 'jax-jaxlib-v%(version)s',
-        'prebuildopts': local_jax_prebuildopts,
-        # Avoid warning (treated as error) in upb/table.c
-        'buildopts': '--bazel_options="--copt=-Wno-maybe-uninitialized"',
+        'start_dir': f'jax-jax-v{version}',
+        # fix jaxlib version - removes .dev suffix
+        'prebuildopts': 'export JAXLIB_RELEASE=%(version)s && ',
+        'buildopts': _jaxlib_buildopts,
     }),
 ]
+# BackendsTest is failing on systems with GPU detected when cpu version is built
+_failing_backend_test = 'tests/api_test.py::BackendsTest::test_no_backend_warning_on_cpu_if_platform_specified'
+# Failing after fix jaxlib version by JAXLIB_RELEASE
+_failing_version_tests = 'tests/version_test.py'
+_runtest_cmd = f"pytest -n %(parallel)s tests --deselect={_failing_backend_test} --deselect={_failing_version_tests}"
 
 exts_list = [
-    ('opt_einsum', '3.3.0', {
-        'checksums': ['59f6475f77bbc37dcf7cd748519c0ec60722e91e63ca114e68821c0c54a46549'],
-    }),
-    ('etils', '0.8.0', {
-        'checksums': ['d1d5af7bd9c784a273c4e1eccfaa8feaca5e0481a08717b5313fa231da22a903'],
-    }),
     (name, version, {
-        'source_tmpl': '%(name)s-v%(version)s.tar.gz',
-        'source_urls': ['https://github.com/google/jax/archive/'],
-        'patches': [
-            'jax-0.3.23_relax-testPoly5-tolerance.patch',
-            'jax-0.3.25_skip-qdwh-test-rank-deficient-deficient.patch',
+        "patches": ["jax-0.6.2_version.patch"],
+        "testinstall": True,
+        "runtest": _runtest_cmd,
+        "sources": [
+            {"source_urls": ["https://github.com/google/jax/archive/"], "filename": "%(name)s-v%(version)s.tar.gz"}
         ],
-        'checksums': [
-            {'jax-v0.3.25.tar.gz': '49e8ce88ddd7dd0de86116c9d75d98a577a9061377ec423493fbac5ea29f79f0'},
-            {'jax-0.3.23_relax-testPoly5-tolerance.patch':
-             'be64bf36dde4884a97b6c8bb22c6b14ab5b24033cd40bfe7ce18363c55c30e87'},
-            {'jax-0.3.25_skip-qdwh-test-rank-deficient-deficient.patch':
-             '70f16f2dba03ab162ce6e13ea61774524b485e9630209bbd4bec81fd16c8812f'},
+        "checksums": [
+            {"jax-v0.6.2.tar.gz": "d46cb98795f2c1ccdf2b081e02d9d74b659063679a80beb001ad17d482a60e17"},
+            {"jax-0.6.2_version.patch": "e15615fd9f4e1698f7c5fd384f146d7b2dbfde3d4657b69bd2d044d75c9fb1d4"},
         ],
-        'runtest': "pytest -n %(parallel)s tests",
     }),
 ]
 
-moduleclass = 'tools'
+moduleclass = 'ai'

pavelToman · 2025-07-29T18:51:20Z

@boegelbot please test @ jsc-zen3 EB_ARGS="--include-easyblocks-from-pr 3852"

boegelbot · 2025-07-29T19:00:11Z

@pavelToman: Request for testing this PR well received on jsczen3l1.int.jsc-zen3.fz-juelich.de

PR test command 'if [[ develop != 'develop' ]]; then EB_BRANCH=develop ./easybuild_develop.sh 2> /dev/null 1>&2; EB_PREFIX=/home/boegelbot/easybuild/develop source init_env_easybuild_develop.sh; fi; EB_PR=23526 EB_ARGS="--include-easyblocks-from-pr 3852" EB_CONTAINER= EB_REPO=easybuild-easyconfigs EB_BRANCH=develop /opt/software/slurm/bin/sbatch --job-name test_PR_23526 --ntasks=8 ~/boegelbot/eb_from_pr_upload_jsc-zen3.sh' executed!

exit code: 0
output:

Submitted batch job 7393

Test results coming soon (I hope)...

Details

- notification for comment with ID 3133665857 processed

Message to humans: this is just bookkeeping information for me,
it is of no use to you (unless you think I have a bug, which I don't).

pavelToman · 2025-07-30T11:35:11Z

Test report by @pavelToman
SUCCESS
Build succeeded for 1 out of 1 (1 easyconfigs in total)
node3908.accelgor.os - Linux RHEL 9.4, x86_64, AMD EPYC 7413 24-Core Processor, Python 3.9.18
See https://gist.github.com/pavelToman/d514ec6e2700dd94c9fb085996423cda for a full test report.

pavelToman · 2025-07-30T17:20:26Z

Test report by @pavelToman
SUCCESS
Build succeeded for 1 out of 1 (1 easyconfigs in total)
node4302.litleo.os - Linux RHEL 9.4, x86_64, AMD EPYC 9454P 48-Core Processor, Python 3.9.18
See https://gist.github.com/pavelToman/b7e85461f222b38aa0ef5b020b487975 for a full test report.

pavelToman · 2025-07-30T18:06:59Z

Test report by @pavelToman
FAILED
Build succeeded for 0 out of 1 (1 easyconfigs in total)
node4205.shinx.os - Linux RHEL 9.4, x86_64, AMD EPYC 9654 96-Core Processor, Python 3.9.18
See https://gist.github.com/pavelToman/f49fee5df51734507b70f986c03d9f92 for a full test report.

pavelToman · 2025-07-31T09:50:35Z

Test report by @pavelToman FAILED Build succeeded for 0 out of 1 (1 easyconfigs in total) node4205.shinx.os - Linux RHEL 9.4, x86_64, AMD EPYC 9654 96-Core Processor, Python 3.9.18 See https://gist.github.com/pavelToman/f49fee5df51734507b70f986c03d9f92 for a full test report.

Failed with 8 cpus + 40gb memory.
Passed with 50gb memory

pavelToman · 2025-07-31T09:58:43Z

v0.6.2 will be for GCC 13.3.0 based toolchain: #23544
Closing this in favor of: #23530

adding easyconfigs: jax-0.6.2-gfbf-2025a.eb and patches: jax-0.6.2_ve…

9a580cf

…rsion.patch

pavelToman added the update label Jul 29, 2025

pavelToman mentioned this pull request Jul 29, 2025

updated of Python packages on top of Python 3.13 for courses (Py4Sci) vscentrum/vsc-software-stack#578

Closed

25 tasks

pavelToman mentioned this pull request Jul 29, 2025

{ai}[gfbf/2024a] jax v0.4.35, ml_dtypes v0.5.0 w/ CUDA 12.6.0 WIP #21924

Draft

9 tasks

Update jax-0.6.2-gfbf-2025a.eb - typo

8c2e43f

pavelToman mentioned this pull request Jul 30, 2025

{ai}[gfbf/2025a] jax v0.7.0 #23530

Merged

pavelToman closed this Jul 31, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

{ai}[gfbf/2025a] jax v0.6.2#23526

{ai}[gfbf/2025a] jax v0.6.2#23526
pavelToman wants to merge 2 commits intoeasybuilders:developfrom
pavelToman:20250729204406_new_pr_jax062

pavelToman commented Jul 29, 2025 •

edited

Loading

Uh oh!

github-actions bot commented Jul 29, 2025 •

edited

Loading

Uh oh!

pavelToman commented Jul 29, 2025

Uh oh!

boegelbot commented Jul 29, 2025

Uh oh!

pavelToman commented Jul 30, 2025

Uh oh!

pavelToman commented Jul 30, 2025

Uh oh!

pavelToman commented Jul 30, 2025

Uh oh!

pavelToman commented Jul 31, 2025

Uh oh!

pavelToman commented Jul 31, 2025 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

pavelToman commented Jul 29, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions bot commented Jul 29, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Updated software jax-0.6.2-gfbf-2025a.eb

Uh oh!

pavelToman commented Jul 29, 2025

Uh oh!

boegelbot commented Jul 29, 2025

Uh oh!

pavelToman commented Jul 30, 2025

Uh oh!

pavelToman commented Jul 30, 2025

Uh oh!

pavelToman commented Jul 30, 2025

Uh oh!

pavelToman commented Jul 31, 2025

Uh oh!

pavelToman commented Jul 31, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

pavelToman commented Jul 29, 2025 •

edited

Loading

github-actions bot commented Jul 29, 2025 •

edited

Loading

Updated software `jax-0.6.2-gfbf-2025a.eb`

pavelToman commented Jul 31, 2025 •

edited

Loading