Add py-jax{,lib} from spack upstream (0.19.0)

packages/py-jax/package.py

+# Copyright 2013-2021 Lawrence Livermore National Security, LLC and other
+# Spack Project Developers. See the top-level COPYRIGHT file for details.
+#
+# SPDX-License-Identifier: (Apache-2.0 OR MIT)
+
+
+from spack.package import *
+
+
+class PyJax(PythonPackage, CudaPackage):
+    """JAX is Autograd and XLA, brought together for high-performance
+    machine learning research. With its updated version of Autograd,
+    JAX can automatically differentiate native Python and NumPy
+    functions. It can differentiate through loops, branches,
+    recursion, and closures, and it can take derivatives of
+    derivatives of derivatives. It supports reverse-mode
+    differentiation (a.k.a. backpropagation) via grad as well as
+    forward-mode differentiation, and the two can be composed
+    arbitrarily to any order."""
+
+    homepage = "https://github.com/google/jax"
+    pypi = "jax/jax-0.2.25.tar.gz"
+
+    version("0.3.25", sha256="18bea69321cb95ea5ea913adfe5e2c1d453cade9d4cfd0dc814ecba9fc0cb6e3")
+    version("0.2.25", sha256="822e8d1e06257eaa0fdc4c0a0686c4556e9f33647fa2a766755f984786ae7446")
+
+    variant("cuda", default=True, description="CUDA support")
+
+    depends_on("python@3.7:", type=("build", "run"))
+    depends_on("py-setuptools", type="build")
+    depends_on("py-numpy@1.18:", type=("build", "run"))
+    depends_on("py-absl-py", type=("build", "run"))
+    depends_on("py-opt-einsum", type=("build", "run"))
+    depends_on("py-scipy@1.2.1:", type=("build", "run"))
+    depends_on("py-typing-extensions", type=("build", "run"))
+    depends_on("py-jaxlib@0.1.69:", type=("build", "run"), when="~cuda")
+    depends_on("py-jaxlib@0.1.69:+cuda", type=("build", "run"), when="+cuda")
+    depends_on("py-jaxlib@0.3.22:", type=("build", "run"), when="@0.3.25:~cuda")
+    depends_on("py-jaxlib@0.3.22:+cuda", type=("build", "run"), when="@0.3.25:+cuda")
+    for arch in CudaPackage.cuda_arch_values:
+        depends_on(
+            "py-jaxlib@0.1.69:+cuda cuda_arch={0}".format(arch),
+            type=("build", "run"),
+            when="cuda_arch={0}".format(arch),
+        )
+        depends_on(
+            "py-jaxlib@0.3.22:+cuda cuda_arch={0}".format(arch),
+            type=("build", "run"),
+            when="@0.3.25:cuda_arch={0}".format(arch),
+        )

packages/py-jaxlib/package.py

+# Copyright 2013-2021 Lawrence Livermore National Security, LLC and other
+# Spack Project Developers. See the top-level COPYRIGHT file for details.
+#
+# SPDX-License-Identifier: (Apache-2.0 OR MIT)
+
+import tempfile
+
+import llnl.util.tty as tty
+
+from spack.package import *
+
+
+class PyJaxlib(PythonPackage, CudaPackage):
+    """XLA library for Jax"""
+
+    homepage = "https://github.com/google/jax"
+    url = "https://github.com/google/jax/archive/refs/tags/jaxlib-v0.1.74.tar.gz"
+
+    tmp_path = ""
+
+    version("0.3.25", sha256="73ebc7868631cd9d520385557bbd7f08762d748a5a6a1bebef0f3b8d7ba748ef")
+    version("0.1.74", sha256="bbc78c7a4927012dcb1b7cd135c7521f782d7dad516a2401b56d3190f81afe35")
+
+    depends_on("python@3.7:", type=("build", "run"))
+    depends_on("py-setuptools", type="build")
+    depends_on("py-numpy@1.18:", type=("build", "run"))
+    depends_on("py-numpy@1.20:", type=("build", "run"), when="@0.3.16:")
+    depends_on("py-opt-einsum", type=("build", "run"), when="@0.2.8:")
+    depends_on("py-scipy", type=("build", "run"))
+    depends_on("py-scipy@1.5:", type=("build", "run"), when="@0.3.14:")
+    depends_on("py-typing-extensions", type=("build", "run"), when="@0.2.23:")
+    depends_on("py-absl-py", type=("build", "run"))
+    depends_on("py-flatbuffers@1.12:2", type=("build", "run"))
+    # Bazel 5 not yet supported: https://github.com/google/jax/issues/8440
+    depends_on("bazel@4.1.0:4", type=("build"), when="@:0.3.5")
+    depends_on("bazel@5.1.1:", type=("build"), when="@0.3.6:")
+    depends_on("cudnn@8.0.5:", when="+cuda")
+    depends_on("cuda@11.1:", when="+cuda")
+
+    def install(self, spec, prefix):
+        args = []
+        args.append("build/build.py")
+        if "+cuda" in spec:
+            args.append("--enable_cuda")
+            args.append("--cuda_path={0}".format(self.spec["cuda"].prefix))
+            args.append("--cudnn_path={0}".format(self.spec["cudnn"].prefix))
+            capabilities = ",".join(
+                "{0:.1f}".format(float(i) / 10.0) for i in spec.variants["cuda_arch"].value
+            )
+            args.append("--cuda_compute_capabilities={0}".format(capabilities))
+        args.append("--bazel_startup_options=" "--output_user_root={0}".format(self.buildtmp))
+        python(*args)
+        with working_dir(self.tmp_path):
+            tty.warn("in dir " + self.tmp_path)
+            args = std_pip_args + ["--prefix=" + self.prefix, "."]
+            pip(*args)
+        remove_linked_tree(self.tmp_path)
+        remove_linked_tree(self.buildtmp)
+
+    def patch(self):
+        self.tmp_path = tempfile.mkdtemp(prefix="spack")
+        self.buildtmp = tempfile.mkdtemp(prefix="spack")
+        filter_file(
+            """f"--output_path={output_path}",""",
+            """f"--output_path={output_path}","""
+            """f"--sources_path=%s","""
+            """f"--nohome_rc'","""
+            """f"--nosystem_rc'",""" % self.tmp_path,
+            "build/build.py",
+        )
+        filter_file(
+            "args = parser.parse_args()",
+            "args,junk = parser.parse_known_args()",
+            "build/build_wheel.py",
+            string=True,
+        )
+