dm-haiku-0.0.9-foss-2022a-CUDA-11.7.0.eb

easyblock = 'PythonBundle'

name = 'dm-haiku'
version = '0.0.9'
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': '2022a'}

dependencies = [
    ('CUDA', '11.7.0', '', SYSTEM),
    ('Python', '3.10.4'),
    ('SciPy-bundle', '2022.05'),
    ('jax', '0.3.25', versionsuffix),  # required by jmp, also provides absl-py
]

use_pip = True

exts_list = [
    ('jmp', '0.0.4', {
        'checksums': ['5dfeb0fd7c7a9f72a70fff0aab9d0cbfae32a809c02f4037ff3485ceb33e1730'],
    }),
    (name, version, {
        'modulename': 'haiku',
        'checksums': ['97752b32cdbe5a3e2d1c60ea884d33eb4b36e7d410000f0d61b571417c925435'],
    }),
]

sanity_pip_check = True

moduleclass = 'lib'