Skip to content
Snippets Groups Projects
Commit 7e1d127e authored by Alexandre Strube's avatar Alexandre Strube
Browse files

set XLA_FLAGS for cuda

parent 859811ee
No related branches found
No related tags found
No related merge requests found
Show whitespace changes
Inline Side-by-side
Golden_Repo/j/JAX/JAX-0.2.9-gcccoremkl-9.3.0-2020.2.254-Python-3.8.5.eb
+ 7
− 4
View file @ 7e1d127e
import os
easyblock = 'PythonPackage' easyblock = 'PythonPackage'
name = 'JAX' name = 'JAX'
...@@ -14,6 +15,7 @@ site_contacts = 'a.strube@fz-juelich.de' ...@@ -14,6 +15,7 @@ site_contacts = 'a.strube@fz-juelich.de'
dependencies = [ dependencies = [
('binutils', '2.34'), ('binutils', '2.34'),
('CUDA', '11.0', '', SYSTEM),
('Python', '3.8.5'), ('Python', '3.8.5'),
('SciPy-Stack', '2020', versionsuffix, ('gcccoremkl', '9.3.0-2020.2.254')), ('SciPy-Stack', '2020', versionsuffix, ('gcccoremkl', '9.3.0-2020.2.254')),
('cuDNN', '8.0.2.39', '-CUDA-%s' % local_cudaver, True), ('cuDNN', '8.0.2.39', '-CUDA-%s' % local_cudaver, True),
...@@ -56,10 +58,11 @@ exts_list = [ ...@@ -56,10 +58,11 @@ exts_list = [
}), }),
] ]
# This should be modextravars, but life is unfair
modextravars = { modluafooter = """
'XLA_FLAGS': '--xla_gpu_cuda_data_dir=$EBROOTCUDA', setenv("CUDA_DIR", os.getenv("EBROOTCUDA"))
} setenv("XLA_FLAGS", "--xla_gpu_cuda_data_dir=" .. os.getenv("EBROOTCUDA"))
"""
sanity_check_paths = { sanity_check_paths = {
'files': [], 'files': [],
......
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment