Skip to content

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

Closed
pavelToman wants to merge 2 commits intoeasybuilders:developfrom
pavelToman:20250729204406_new_pr_jax062
Closed

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

Conversation

@pavelToman
Copy link
Collaborator

@pavelToman pavelToman commented Jul 29, 2025

@github-actions
Copy link

github-actions bot commented Jul 29, 2025

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
Copy link
Collaborator Author

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

@boegelbot
Copy link
Collaborator

@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
Copy link
Collaborator Author

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
Copy link
Collaborator Author

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
Copy link
Collaborator Author

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
Copy link
Collaborator Author

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
Copy link
Collaborator Author

pavelToman commented Jul 31, 2025

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

@pavelToman pavelToman closed this Jul 31, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

updated of Python packages on top of Python 3.13 for courses (Py4Sci)

2 participants