Skip to content

Commit 7279372

Browse files
laraPPrlaraPPr
authored
Merge pull request #26540 from JarneRenders/20260727152936_new_pr_Flax0112
{ai,lib}[gfbf/2024a] Flax v0.11.2, Optax v0.2.8 w/ CUDA 12.6.0
2 parents 1f244da + 6ce37b4 commit 7279372

4 files changed

Lines changed: 156 additions & 0 deletions

File tree

‎easybuild/easyconfigs/f/Flax/Flax-0.11.2-gfbf-2024a-CUDA-12.6.0.eb‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
easyblock = 'PythonBundle'
2+
3+
name = 'Flax'
4+
version = '0.11.2'
5+
versionsuffix = '-CUDA-%(cudaver)s'
6+
7+
homepage = 'https://flax.readthedocs.io'
8+
description = """Flax is a high-performance neural network library and ecosystem for JAX that is
9+
designed for flexibility: Try new forms of training by forking an example and
10+
by modifying the training loop, not by adding features to a framework."""
11+
12+
toolchain = {'name': 'gfbf', 'version': '2024a'}
13+
14+
builddependencies = [
15+
('hatchling', '1.27.0'),
16+
('trove-classifiers', '2026.1.14.14'), # needed for hatchling to recognize Python :: 3.15
17+
]
18+
19+
dependencies = [
20+
('CUDA', '12.6.0', '', SYSTEM),
21+
('Python', '3.12.3'),
22+
('SciPy-bundle', '2024.05'),
23+
('jax', '0.6.2', versionsuffix),
24+
('Optax', '0.2.8', versionsuffix),
25+
('protobuf-python', '5.28.0'),
26+
('PyYAML', '6.0.2'),
27+
('tensorstore', '0.1.72'),
28+
]
29+
30+
exts_list = [
31+
('nest_asyncio', '1.6.0', {
32+
'checksums': ['6f172d5449aca15afd6c646851f4e31e02c598d553a667e38cafa997cfec55fe'],
33+
}),
34+
('aiofiles', '25.1.0', { # for orbax_checkpoint
35+
'checksums': ['a8d728f0a29de45dc521f18f07297428d56992a742f0cd2701ba86e44d23d5b2'],
36+
}),
37+
('humanize', '4.16.0', { # for orbax_checkpoint
38+
'preinstallopts': "sed -i '/^license = /d' pyproject.toml && ",
39+
'checksums': ['7dc2244a2f84a4bfb1d36c37bac80cd78e35cdc5c119206d87b018e1445f3a3f'],
40+
}),
41+
('orbax_checkpoint', '0.11.24', {
42+
'modulename': 'orbax.checkpoint',
43+
'checksums': ['4e7afe927d1ed6d8160bacf5ed4fef56c1370320e0ebdfda213c6351a2e3c0d0'],
44+
}),
45+
('treescope', '0.1.10', {
46+
'checksums': ['20f74656f34ab2d8716715013e8163a0da79bdc2554c16d5023172c50d27ea95'],
47+
}),
48+
('flax', version, {
49+
'checksums': ['55452529b70c704128075a17f7a54d94f1f90b0f2d9498bdc9578cae3646e460'],
50+
}),
51+
]
52+
53+
moduleclass = 'ai'

‎easybuild/easyconfigs/f/Flax/Flax-0.11.2-gfbf-2024a.eb‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
easyblock = 'PythonBundle'
2+
3+
name = 'Flax'
4+
version = '0.11.2'
5+
6+
homepage = 'https://flax.readthedocs.io'
7+
description = """Flax is a high-performance neural network library and ecosystem for JAX that is
8+
designed for flexibility: Try new forms of training by forking an example and
9+
by modifying the training loop, not by adding features to a framework."""
10+
11+
toolchain = {'name': 'gfbf', 'version': '2024a'}
12+
13+
builddependencies = [
14+
('hatchling', '1.27.0'),
15+
('trove-classifiers', '2026.1.14.14'), # needed for hatchling to recognize Python :: 3.15
16+
]
17+
18+
dependencies = [
19+
('Python', '3.12.3'),
20+
('SciPy-bundle', '2024.05'),
21+
('jax', '0.6.2'),
22+
('Optax', '0.2.8'),
23+
('protobuf-python', '5.28.0'),
24+
('PyYAML', '6.0.2'),
25+
('tensorstore', '0.1.72'),
26+
]
27+
28+
exts_list = [
29+
('nest_asyncio', '1.6.0', {
30+
'checksums': ['6f172d5449aca15afd6c646851f4e31e02c598d553a667e38cafa997cfec55fe'],
31+
}),
32+
('aiofiles', '25.1.0', { # for orbax_checkpoint
33+
'checksums': ['a8d728f0a29de45dc521f18f07297428d56992a742f0cd2701ba86e44d23d5b2'],
34+
}),
35+
('humanize', '4.16.0', { # for orbax_checkpoint
36+
'preinstallopts': "sed -i '/^license = /d' pyproject.toml && ",
37+
'checksums': ['7dc2244a2f84a4bfb1d36c37bac80cd78e35cdc5c119206d87b018e1445f3a3f'],
38+
}),
39+
('orbax_checkpoint', '0.11.24', {
40+
'modulename': 'orbax.checkpoint',
41+
'checksums': ['4e7afe927d1ed6d8160bacf5ed4fef56c1370320e0ebdfda213c6351a2e3c0d0'],
42+
}),
43+
('treescope', '0.1.10', {
44+
'checksums': ['20f74656f34ab2d8716715013e8163a0da79bdc2554c16d5023172c50d27ea95'],
45+
}),
46+
('flax', version, {
47+
'checksums': ['55452529b70c704128075a17f7a54d94f1f90b0f2d9498bdc9578cae3646e460'],
48+
}),
49+
]
50+
51+
moduleclass = 'ai'

‎easybuild/easyconfigs/o/Optax/Optax-0.2.8-gfbf-2024a-CUDA-12.6.0.eb‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
easyblock = 'PythonBundle'
2+
3+
name = 'Optax'
4+
version = '0.2.8'
5+
versionsuffix = '-CUDA-%(cudaver)s'
6+
7+
homepage = 'https://github.com/deepmind/optax'
8+
description = """Optax is a gradient processing and optimization library for JAX."""
9+
10+
toolchain = {'name': 'gfbf', 'version': '2024a'}
11+
12+
dependencies = [
13+
('CUDA', '12.6.0', '', SYSTEM),
14+
('Python', '3.12.3'),
15+
('SciPy-bundle', '2024.05'),
16+
('jax', '0.6.2', versionsuffix),
17+
]
18+
19+
exts_list = [
20+
('optax', version, {
21+
'checksums': ['5b225b35066fc3eebaa4d798f1b4173b4d57d1a480610908981f8343b50af0b0'],
22+
}),
23+
]
24+
25+
sanity_check_commands = ["python -c 'from optax import GradientTransformation'"]
26+
27+
moduleclass = 'lib'

‎easybuild/easyconfigs/o/Optax/Optax-0.2.8-gfbf-2024a.eb‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
easyblock = 'PythonBundle'
2+
3+
name = 'Optax'
4+
version = '0.2.8'
5+
6+
homepage = 'https://github.com/deepmind/optax'
7+
description = """Optax is a gradient processing and optimization library for JAX."""
8+
9+
toolchain = {'name': 'gfbf', 'version': '2024a'}
10+
11+
dependencies = [
12+
('Python', '3.12.3'),
13+
('SciPy-bundle', '2024.05'),
14+
('jax', '0.6.2'),
15+
]
16+
17+
exts_list = [
18+
('optax', version, {
19+
'checksums': ['5b225b35066fc3eebaa4d798f1b4173b4d57d1a480610908981f8343b50af0b0'],
20+
}),
21+
]
22+
23+
sanity_check_commands = ["python -c 'from optax import GradientTransformation'"]
24+
25+
moduleclass = 'lib'

0 commit comments

Comments
 (0)