{lib}[foss/2024a] dm-haiku v0.0.13 w/ CUDA 12.6.0

easybuild/easyconfigs/d/dm-haiku/dm-haiku-0.0.13-foss-2024a-CUDA-12.6.0.eb

@@ -0,0 +1,56 @@
+# update 0.0.12: Thomas Hoffmann (EMBL) 
+easyblock = 'PythonBundle'
+
+name = 'dm-haiku'
+version = '0.0.13'
+versionsuffix = '-CUDA-%(cudaver)s'
+
+homepage = 'https://github.com/deepmind/dm-haiku'
+description = """Haiku is a simple neural network library for JAX developed by some of the authors of Sonnet, a neural
+network library for TensorFlow."""
+
+toolchain = {'name': 'foss', 'version': '2024a'}
+
+builddependencies = [
+    ('poetry', '1.8.3')
+]
+
+dependencies = [
+    ('Python', '3.12.3'),
+    ('SciPy-bundle', '2024.05'),
+    ('jax', '0.4.34', versionsuffix),  # required by jmp, also provides absl-py
+    ('PyYAML', '6.0.2'),
+    ('CUDA', '12.6.0', '', SYSTEM),
+    ('tensorstore', '0.1.72'),
+    ('protobuf-python', '5.28.0'),
+    ('Optax', '0.2.4', versionsuffix),
+]
+
+use_pip = True
+
+exts_list = [
+    ('jmp', '0.0.4', {
+        'checksums': ['5dfeb0fd7c7a9f72a70fff0aab9d0cbfae32a809c02f4037ff3485ceb33e1730'],
+    }),
+    ('flax', '0.10.2', {
+        'checksums': ['6f831350026ad48182ba6588bb4dd72dc1084985d9aca923254cb3e4c78d75f3'],
+    }),
+    ('nest_asyncio', '1.6.0', {
+        'checksums': ['6f172d5449aca15afd6c646851f4e31e02c598d553a667e38cafa997cfec55fe'],
+    }),
+    ('humanize', '4.11.0', {
+        'checksums': ['e66f36020a2d5a974c504bd2555cf770621dbdbb6d82f94a6857c0b1ea2608be'],
+    }),
+    ('orbax_checkpoint', '0.10.2', {
+        'modulename': 'orbax.checkpoint',
+        'checksums': ['e575ebe1f94e5cb6353ab8c9df81de0ca7cddc118645c3bfc17b8344f19d42f1'],
+    }),
+    ('dm_haiku', version, {
+        'modulename': 'haiku',
+        'checksums': ['029bb91b5b1edb0d3fe23304d3bf12a545ea6e485041f7f5d8c8d85ebcf6e17d'],
+    }),
+]
+
+sanity_pip_check = True
+
+moduleclass = 'lib'

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

@@ -0,0 +1,37 @@
+easyblock = 'PythonBundle'
+
+name = 'Optax'
+version = '0.2.4'
+versionsuffix = '-CUDA-%(cudaver)s'
+
+homepage = 'https://github.com/deepmind/optax'
+description = """Optax is a gradient processing and optimization library for JAX."""
+
+toolchain = {'name': 'gfbf', 'version': '2024a'}
+
+dependencies = [
+    ('CUDA', '12.6.0', '', SYSTEM),
+    ('Python', '3.12.3'),
+    ('SciPy-bundle', '2024.05'),
+    ('jax', '0.4.34', versionsuffix),
+]
+
+use_pip = True
+
+exts_list = [
+    ('chex', '0.1.87', {
+        'checksums': ['0096d89cc8d898bb521ef4bfbf5c24549022b0e5b301f529ab57238896fe6c5d'],
+    }),
+    ('toolz', '1.0.0', {
+        'checksums': ['2c86e3d9a04798ac556793bced838816296a2f085017664e4995cb40a1047a02'],
+    }),
+    ('optax', version, {
+        'checksums': ['4e05d3d5307e6dde4c319187ae36e6cd3a0c035d4ed25e9e992449a304f47336'],
+    }),
+]
+
+sanity_pip_check = True
+
+sanity_check_commands = ["python -c 'from optax import GradientTransformation'"]
+
+moduleclass = 'lib'