{tools}[foss/2020b] jax v0.2.19

easybuild/easyconfigs/j/jax/jax-0.2.18-correct_libraries.patch

@@ -0,0 +1,63 @@
+diff -ruN jax_orig/build/build.py jax/build/build.py
+--- jax_orig/build/build.py	2021-07-30 10:18:43.599233082 +0200
++++ jax/build/build.py	2021-08-02 15:58:09.654165000 +0200
+@@ -41,7 +41,7 @@
+ else:
+   from distutils.spawn import find_executable as which
+ # pylint: enable=g-import-not-at-top
+-
++use_default_shell_env = True
+
+ def is_windows():
+     return sys.platform.startswith("win32")
+Binary files jax_orig/.git/index and jax/.git/index differ
+diff -ruN jax_orig/jaxlib/jax.bzl jax/jaxlib/jax.bzl
+--- jax_orig/jaxlib/jax.bzl	2021-07-30 10:18:47.211049000 +0200
++++ jax/jaxlib/jax.bzl	2021-08-02 10:39:41.567268000 +0200
+@@ -30,3 +30,4 @@
+ if_rocm_is_configured = _if_rocm_is_configured
+ flatbuffer_cc_library = _flatbuffer_cc_library
+ flatbuffer_py_library = _flatbuffer_py_library
++use_default_shell_env = True
+diff -ruN jax_orig/jaxlib/lapack.pyx jax/jaxlib/lapack.pyx
+--- jax_orig/jaxlib/lapack.pyx	2021-07-30 10:18:47.214725015 +0200
++++ jax/jaxlib/lapack.pyx	2021-07-30 10:23:25.084607000 +0200
+@@ -49,6 +49,7 @@
+
+ cdef int _int32_max = 0x7FFFFFFF;
+
++use_default_shell_env = True
+
+ cdef register_cpu_custom_call_target(fn_name, void* fn):
+   cdef const char* name = "xla._CUSTOM_CALL_TARGET"
+diff -ruN jax_orig/WORKSPACE jax/WORKSPACE
+--- jax_orig/WORKSPACE	2021-07-30 10:18:43.503078000 +0200
++++ jax/WORKSPACE	2021-08-03 18:13:58.627602000 +0200
+@@ -7,17 +7,24 @@
+ #    and update the sha256 with the result.
+ http_archive(
+     name = "org_tensorflow",
+-    sha256 = "de22145c9dc51b3ed9399dbee2ab681f094f09e5213d6819451da500f163a14b",
++    # sha256 = "de22145c9dc51b3ed9399dbee2ab681f094f09e5213d6819451da500f163a14b",
+     strip_prefix = "tensorflow-8cc3ffa8d8e4dd659c1534849cf5984ef4ec3532",
+     urls = [
+         "https://github.com/tensorflow/tensorflow/archive/8cc3ffa8d8e4dd659c1534849cf5984ef4ec3532.tar.gz",
+     ],
++    patch_cmds=['sed -i -e "s|arguments = \[args\],|arguments = \[args\],\\n\\tuse_default_shell_env = True,|" third_party/mlir/tblgen.bzl']
+ )
+
+-# For development, one can use a local TF repository instead.
++# # For development, one can use a local TF repository instead.
+ # local_repository(
+ #    name = "org_tensorflow",
+-#    path = "tensorflow",
++#    path = "/apps/gent/CO7/skylake-ib/software/TensorFlow/2.5.0-foss-2020b",
++# )
++# new_local_repository(
++#     name = "org_tensorflow",
++#     path = "/apps/gent/CO7/skylake-ib/software/TensorFlow/2.5.0-foss-2020b/lib/python3.8/site-packages/tensorflow/",
++#     # build_file = "org_tensorflow.BUILD",
++
+ # )
+
+ load("//third_party/pocketfft:workspace.bzl", pocketfft = "repo")

easybuild/easyconfigs/j/jax/jax-0.2.18-foss-2020b.eb

@@ -0,0 +1,35 @@
+easyblock = 'PythonPackage'
+
+name = 'jax'
+version = '0.2.18'
+
+homepage = 'https://pypi.python.org/pypi/jax'
+description = """Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and
+more"""
+
+toolchain = {'name': 'foss', 'version': '2020b'}
+
+# source_urls = [PYPI_SOURCE]
+# sources = [SOURCE_TAR_GZ]
+
+source_urls = ['https://github.com/google/jax/archive/']
+sources = ['%(name)s-v%(version)s.tar.gz']
+patches = [
+    'jax-0.2.18-correct_libraries.patch',
+]
+
+preinstallopts = 'python3 build/build.py --output_user_root=%(builddir)s --bazel_path="$EBROOTBAZEL/bin/bazel" && '
+preinstallopts += 'pip install dist/*.whl && '
+
+dependencies = [
+    ('Python', '3.8.6'),
+    ('Bazel', '3.7.2'),
+    ('SciPy-bundle', '2020.11'),
+    ('Cython', '0.29.22'),
+    # ('TensorFlow', '2.5.0'), do we need this?
+]
+
+use_pip = True
+
+
+moduleclass = 'tools'