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

easybuild/easyconfigs/j/jax/TensorFlow-4039feeb743bc42cd0a3d8146ce63fc05d23eb8d_mlir-build-env.patch

@@ -0,0 +1,12 @@
+make TensorFlow pick up on build environment defined by EasyBuild
+author: Denis Kristak, Pavel Grochal (INUITS)
+--- tensorflow-4039feeb743bc42cd0a3d8146ce63fc05d23eb8d/third_party/mlir/tblgen.bzl.orig	2021-08-13 16:18:12.735593126 +0200
++++ tensorflow-4039feeb743bc42cd0a3d8146ce63fc05d23eb8d/third_party/mlir/tblgen.bzl	2021-08-13 16:18:48.615675683 +0200
+@@ -173,6 +173,7 @@
+         inputs = trans_srcs,
+         executable = ctx.executable.tblgen,
+         arguments = [args],
++        use_default_shell_env = True,
+         mnemonic = "TdGenerate",
+     )
+

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

@@ -0,0 +1,113 @@
+# This file is an EasyBuild reciPY as per https://github.com/easybuilders/easybuild
+# Author: Denis Kristak
+easyblock = 'PythonBundle'
+
+name = 'jax'
+version = '0.2.19'
+local_jaxlib_ver = '0.1.70'
+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'}
+
+builddependencies = [
+    ('Bazel', '3.7.2'),
+    ('pytest-xdist', '2.3.0'),
+    ('pytest-benchmark', '3.4.1'),
+]
+
+dependencies = [
+    ('Python', '3.8.6'),
+    ('SciPy-bundle', '2020.11'),
+    ('Cython', '0.29.22'),
+    ('flatbuffers-python', '1.12'),
+]
+
+# downloading TensorFlow tarball to avoid that Bazel downloads it during the build
+# note: this *must* be the exact same commit as used in jaxlib-*/WORKSPACE
+local_tf_commit = '4039feeb743bc42cd0a3d8146ce63fc05d23eb8d'
+
+local_tf_builddir = "%%(builddir)s/tensorflow-%s" % local_tf_commit
+
+local_jax_prebuildopts = "sed -i -e 's$pathToSed$%s$g' WORKSPACE && " % local_tf_builddir
+
+local_jax_build_cmd = 'python build/build.py '
+local_jax_build_cmd += '--target_cpu_features=native '
+local_jax_build_cmd += '--bazel_startup_options="--output_user_root=%(builddir)s" '
+local_jax_build_cmd += '--bazel_path="$EBROOTBAZEL/bin/bazel" '
+local_jax_build_cmd += '--bazel_options=--subcommands '
+local_jax_build_cmd += '--bazel_options=--jobs=%(parallel)s --bazel_options=--action_env=PYTHONPATH '
+local_jax_build_cmd += '--bazel_options=--action_env=EBPYTHONPREFIXES'
+
+use_pip = True
+
+default_component_specs = {
+    'sources': [SOURCE_TAR_GZ],
+    'source_urls': [PYPI_SOURCE],
+    'start_dir': '%(name)s-%(version)s',
+}
+
+components = [
+    ('absl-py', '0.13.0', {
+        'options': {'modulename': 'absl'},
+        'easyblock': 'PythonPackage',
+        'checksums': ['6953272383486044699fd0e9f00aad167a27e08ce19aae66c6c4b10e7e767793'],
+        'use_pip': True,
+    }),
+    ('jaxlib', local_jaxlib_ver, {
+        'easyblock': 'PythonPackage',
+        'sources': [
+            'jaxlib-v%s.zip' % local_jaxlib_ver,
+            {
+                'download_filename': '%s.tar.gz' % local_tf_commit,
+                'filename': 'tensorflow-%s.tar.gz' % local_tf_commit,
+            }
+        ],
+        'source_urls': [
+            'https://github.com/google/jax/archive/',
+            'https://github.com/tensorflow/tensorflow/archive/'
+        ],
+        'patches': [
+            ('jaxlib-%s_no-tensorflow-download.patch' % local_jaxlib_ver, 1),
+            ('TensorFlow-%s_mlir-build-env.patch' % local_tf_commit, '../tensorflow-%s' % local_tf_commit),
+        ],
+        'checksums': [
+            # jaxlib-v0.1.70.zip
+            '1205b180de70c98899d842f75bb589599c10a3c7a1c6099475c3a7d39b2e44ac',
+            # tensorflow-4039feeb743bc42cd0a3d8146ce63fc05d23eb8d.tar.gz
+            '0f13410284b9186e436350e9617b3bed2d65f1dc1a220fd37ad9ef43c2035663',
+            # jaxlib-0.1.70_no-tensorflow-download.patch
+            '1496c6be9336fb563b997f8da83d690a9b3a0a6e29fdd87fac60409300a2fd95',
+            # TensorFlow-4039feeb743bc42cd0a3d8146ce63fc05d23eb8d-build-env.patch
+            'd067836b1d18d66a8c69fb3db311855c8a56c0106d42a11299a1733c385e9c22',
+        ],
+        'start_dir': 'jax-jaxlib-v%s' % local_jaxlib_ver,
+        'prebuildopts': local_jax_prebuildopts,
+        'use_pip': True,
+        'sanity_pip_check': True,
+        'download_dep_fail': True,
+        'buildcmd': local_jax_build_cmd,
+        'install_src': 'dist/*.whl',
+    }),
+]
+
+exts_list = [
+    ('opt_einsum', '3.3.0', {
+        'checksums': ['59f6475f77bbc37dcf7cd748519c0ec60722e91e63ca114e68821c0c54a46549'],
+    }),
+    (name, version, {
+        'checksums': ['e55534680a05b6ced92dadde307bd3b2207efd1aaaec304e96197b00ad7560d3'],
+    }),
+]
+
+# can't run tests in sanity check step, because build dependencies must be available
+postinstallcmds = [
+    "cd %%(builddir)s/jax-jaxlib-v%s && "
+    "PYTHONPATH=%%(installdir)s/lib/python%%(pyshortver)s/site-packages:$PYTHONPATH "
+    "pytest -n auto tests" % local_jaxlib_ver,
+]
+
+sanity_pip_check = True
+
+moduleclass = 'tools'

easybuild/easyconfigs/j/jax/jaxlib-0.1.70_no-tensorflow-download.patch

@@ -0,0 +1,38 @@
+avoid downloading TensorFlow during the build, add placeholder to inject location of pre-downloaded TensorFlow sources
+author: Denis Kristak (INUITS)
+diff -ruN jax-jaxlib-v0.1.70_orig/WORKSPACE jax-jaxlib-v0.1.70/WORKSPACE
+--- jax-jaxlib-v0.1.70_orig/WORKSPACE	2021-08-12 13:12:54.025608000 +0200
++++ jax-jaxlib-v0.1.70/WORKSPACE	2021-08-12 13:32:51.581961000 +0200
+@@ -5,21 +5,20 @@
+ # b) get the sha256 hash of the commit by running:
+ #    curl -L https://github.com/tensorflow/tensorflow/archive/<git hash>.tar.gz | sha256sum
+ #    and update the sha256 with the result.
+-http_archive(
++# http_archive(
++#     name = "org_tensorflow",
++#     sha256 = "0f13410284b9186e436350e9617b3bed2d65f1dc1a220fd37ad9ef43c2035663",
++#     strip_prefix = "tensorflow-4039feeb743bc42cd0a3d8146ce63fc05d23eb8d",
++#     urls = [
++#         "https://github.com/tensorflow/tensorflow/archive/4039feeb743bc42cd0a3d8146ce63fc05d23eb8d.tar.gz",
++#     ],
++# )
++
++local_repository(
+     name = "org_tensorflow",
+-    sha256 = "0f13410284b9186e436350e9617b3bed2d65f1dc1a220fd37ad9ef43c2035663",
+-    strip_prefix = "tensorflow-4039feeb743bc42cd0a3d8146ce63fc05d23eb8d",
+-    urls = [
+-        "https://github.com/tensorflow/tensorflow/archive/4039feeb743bc42cd0a3d8146ce63fc05d23eb8d.tar.gz",
+-    ],
++    path = "pathToSed",
+ )
+
+-# For development, one can use a local TF repository instead.
+-# local_repository(
+-#    name = "org_tensorflow",
+-#    path = "tensorflow",
+-# )
+-
+ load("//third_party/pocketfft:workspace.bzl", pocketfft = "repo")
+ pocketfft()
+

easybuild/easyconfigs/p/pytest-benchmark/pytest-benchmark-3.4.1-GCCcore-10.2.0.eb

@@ -0,0 +1,41 @@
+easyblock = 'PythonBundle'
+
+name = 'pytest-benchmark'
+version = '3.4.1'
+
+homepage = 'https://github.com/ionelmc/pytest-benchmark'
+description = "A pytest fixture for benchmarking code."
+
+toolchain = {'name': 'GCCcore', 'version': '10.2.0'}
+
+builddependencies = [('binutils', '2.35')]
+
+dependencies = [
+    ('Python', '3.8.6'),
+]
+
+use_pip = True
+
+exts_list = [
+    ('py-cpuinfo', '8.0.0', {
+        'checksums': ['5f269be0e08e33fd959de96b34cd4aeeeacac014dd8305f70eb28d06de2345c5'],
+        'modulename': 'cpuinfo',
+    }),
+    (name, version, {
+        'checksums': ['40e263f912de5a81d891619032983557d62a3d85843f9a9f30b98baea0cd7b47'],
+    }),
+]
+
+sanity_check_paths = {
+    'files': ['bin/cpuinfo', 'bin/pytest-benchmark'],
+    'dirs': ['lib/python%(pyshortver)s/site-packages'],
+}
+
+sanity_check_commands = [
+    "cpuinfo --help",
+    "pytest-benchmark --help",
+]
+
+sanity_pip_check = True
+
+moduleclass = 'tools'

easybuild/easyconfigs/p/pytest-xdist/pytest-xdist-2.3.0-GCCcore-10.2.0.eb

@@ -0,0 +1,62 @@
+easyblock = 'PythonBundle'
+
+name = 'pytest-xdist'
+version = '2.3.0'
+
+homepage = 'https://github.com/pytest-dev/pytest-xdist'
+description = """xdist: pytest distributed testing plugin
+
+The pytest-xdist plugin extends pytest with some unique test execution modes:
+
+    * test run parallelization: if you have multiple CPUs or hosts you
+      can use those for a combined test run. This allows to speed up
+      development or to use special resources of remote machines.
+
+    * --looponfail: run your tests repeatedly in a subprocess. After
+        each run pytest waits until a file in your project changes and
+        then re-runs the previously failing tests. This is repeated
+        until all tests pass after which again a full run is
+        performed.
+
+    * Multi-Platform coverage: you can specify different Python
+      interpreters or different platforms and run tests in parallel on
+      all of them.
+
+Before running tests remotely, pytest efficiently “rsyncs” your
+program source code to the remote place. All test results are reported
+back and displayed to your local terminal. You may specify different
+Python versions and interpreters."""
+
+toolchain = {'name': 'GCCcore', 'version': '10.2.0'}
+
+builddependencies = [('binutils', '2.35')]
+
+dependencies = [
+    ('Python', '3.8.6'),
+]
+
+use_pip = True
+sanity_pip_check = True
+
+exts_list = [
+    ('apipkg', '1.5', {
+        'checksums': ['37228cda29411948b422fae072f57e31d3396d2ee1c9783775980ee9c9990af6'],
+    }),
+    ('execnet', '1.9.0', {
+        'checksums': ['8f694f3ba9cc92cab508b152dcfe322153975c29bda272e2fd7f3f00f36e47c5'],
+    }),
+    ('pytest-forked', '1.3.0', {
+        'checksums': ['6aa9ac7e00ad1a539c41bec6d21011332de671e938c7637378ec9710204e37ca'],
+    }),
+    (name, version, {
+        'modulename': 'xdist',
+        'checksums': ['e8ecde2f85d88fbcadb7d28cb33da0fa29bca5cf7d5967fa89fc0e97e5299ea5'],
+    }),
+]
+
+sanity_check_paths = {
+    'files': [],
+    'dirs': ['lib/python%(pyshortver)s/site-packages'],
+}
+
+moduleclass = 'tools'