APEBench and its dependencies + Updates to equinox module for supporting cuda as well as latest versions (spack#1662)

abhishek1297 · web-flow · commit ee3b10360d54 · 2025-10-14T10:46:46.000+01:00
diff --git a/repos/spack_repo/builtin/packages/py_apebench/package.py b/repos/spack_repo/builtin/packages/py_apebench/package.py
@@ -0,0 +1,63 @@
+# Copyright Spack Project Developers. See COPYRIGHT file for details.
+#
+# SPDX-License-Identifier: (Apache-2.0 OR MIT)
+
+
+from spack_repo.builtin.build_systems.cuda import CudaPackage
+from spack_repo.builtin.build_systems.python import PythonPackage
+
+from spack.package import *
+
+
+class PyApebench(PythonPackage, CudaPackage):
+    """APEBench is a JAX-based tool to evaluate autoregressive neural emulators for PDEs
+    on periodic domains in 1d, 2d, and 3d. It comes with an efficient reference simulator
+    based on spectral methods that is used for procedural data generation (no need to download
+    large datasets with APEBench). Since this simulator can also be embedded into emulator training
+    (e.g., for a "solver-in-the-loop" correction setting), this is the first benchmark suite to
+    support differentiable physics."""
+
+    homepage = "https://tum-pbs.github.io/apebench-paper/"
+    pypi = "apebench/apebench-0.1.1.tar.gz"
+
+    maintainers("abhishek1297")
+    license("MIT", checked_by="abhishek1297")
+
+    version("0.1.1", sha256="c5ddd47799f0799b2c2e72c27d3d81993f6fa218a04b1df93d4c1850e4893bf9")
+
+    depends_on("c", type="build")
+    depends_on("cxx", type="build")
+    depends_on("py-setuptools", type="build")
+    depends_on("python@3.10:3.12", type=("build", "run"))
+
+    with default_args(type="run"):
+        for arch in CudaPackage.cuda_arch_values:
+            cuda_specs = f"+cuda cuda_arch={arch}"
+            with when(cuda_specs):
+                depends_on(f"py-jaxlib {cuda_specs}")
+                depends_on(f"py-equinox@0.11.3: {cuda_specs}")
+                depends_on(f"py-exponax@0.1.0 {cuda_specs}")
+                depends_on(f"py-pdequinox@0.1.2 {cuda_specs}")
+                depends_on(f"py-trainax@0.0.2 {cuda_specs}")
+
+        depends_on("py-jax@0.4.13:")
+        depends_on("py-jaxtyping@0.2.20:")
+        depends_on("py-typing-extensions@4.5.0:")
+        depends_on("py-tqdm@4.63.2:")
+        depends_on("py-matplotlib@3.8.1:")
+        depends_on("py-pandas@2.2.0:")
+        depends_on("py-seaborn@0.13.0:")
+        depends_on("py-optax@0.2.0:")
+
+        with when("~cuda"):
+            depends_on("py-equinox@0.11.3:")
+            depends_on("py-exponax@0.1.0")
+            depends_on("py-pdequinox@0.1.2")
+            depends_on("py-trainax@0.0.2")
+
+    def setup_run_environment(self, env):
+        if "+cuda" in self.spec:
+            cuda_home = self.spec["cuda"].prefix
+            # This is an irrelevant lib path and it is purely used by NVIDIA profilers.
+            # But, since JAX throws RuntimeError on it, we set this path.
+            env.prepend_path("LD_LIBRARY_PATH", f"{cuda_home}/extras/CUPTI/lib64")
diff --git a/repos/spack_repo/builtin/packages/py_equinox/package.py b/repos/spack_repo/builtin/packages/py_equinox/package.py
@@ -2,32 +2,60 @@
 #
 # SPDX-License-Identifier: (Apache-2.0 OR MIT)
 
+from spack_repo.builtin.build_systems.cuda import CudaPackage
 from spack_repo.builtin.build_systems.python import PythonPackage
 
 from spack.package import *
 
 
-class PyEquinox(PythonPackage):
-    """Equinox is a comprehensive JAX library that provides a wide
-    range of tools and features not found in core JAX, including neural networks
-    with PyTorch-like syntax, filtered APIs for transformations, PyTree manipulation
-    routines, and advanced features like runtime errors."""
+class PyEquinox(PythonPackage, CudaPackage):
+    """Equinox is your one-stop [JAX](https://github.com/google/jax) library,
+    for everything you need that isn't already in core JAX:
+    - neural networks (or more generally any model), with easy-to-use PyTorch-like syntax;
+    - filtered APIs for transformations;
+    - useful PyTree manipulation routines;
+    - advanced features like runtime errors;
 
-    homepage = "https://docs.kidger.site/equinox/"
+    and best of all, Equinox isn't a framework: everything you write in Equinox is compatible
+    with anything else in JAX or the ecosystem."""
 
-    pypi = "equinox/equinox-0.11.10.tar.gz"
+    homepage = "https://docs.kidger.site/equinox/"
+    pypi = "equinox/equinox-0.13.1.tar.gz"
 
-    license("Apache-2.0", checked_by="viperML")
+    maintainers("abhishek1297")
+    license("Apache-2.0", checked_by="abhishek1297")
 
+    version("0.13.1", sha256="e90f11cfe66b2f73f5c172260a17c48851794a0f243dd2cbe4ea70f4c90cbd07")
+    version("0.13.0", sha256="d59615be722373e9d66e0ba78462964e6357fb76a8b1b98c2c6027961b778a69")
+    version("0.12.2", sha256="648e4206bbc53b228922e8f18cd3cffe543ddda1172c0002f8954e484bab0023")
+    version("0.12.1", sha256="7ed4b84553cb59d4930185f87ac2c1121aab2b38999be9499c021e7583a7ed0d")
+    version("0.12.0", sha256="6a99877376cfc168cfe44220a734740926bf23eb9c0cd0d7fdc49adfec4d78ca")
     version("0.11.12", sha256="bee22aabaf7ee0cde6f2ae58cf3b981dea73d47e297361a0203e299208ef1739")
+    version("0.11.11", sha256="648072c1384adc3528930a3bf089246fd77aa873310a19f1f21c08e7681f95a7")
     version("0.11.10", sha256="f3e7d5545b71e427859a28050526d09adb6b20285c47476a606328a0b96c9509")
+    version("0.11.9", sha256="e0f0fa5ea597949492d201ab4d08b05c2d5b4020c65a1778aedf6ad76c2c4fe7")
+    version("0.11.8", sha256="d1e91a05e41bb9538db72a8e15d26daf958348c26714533434c88c5ec0c0b0ef")
+    version("0.11.7", sha256="96e0216a9d822ec4b1465b0cbfbab14a36fb7e7d62c55f521287db3aaaa251be")
+    version("0.11.6", sha256="e237c25e446960ed479f086df240d4dd779bb0917bafc76811d341ccac76b712")
+    version("0.11.5", sha256="5e0ca252eeb20cc5eece225d3d35137c7e57f998a1c6422a1972db9e5c68b7f6")
+    version("0.11.4", sha256="0033d9731083f402a855b12a0777a80aa8507651f7aa86d9f0f9503bcddfd320")
+    version("0.11.3", sha256="a1273cc28c60d3131ac596f8a0f5c7dd384729e6cddae86e7be05f026880e8e0")
 
     depends_on("py-hatchling", type="build")
+    depends_on("python@3.10:3.12", type=("build", "run"))
+
+    with default_args(type="run"):
+        for arch in CudaPackage.cuda_arch_values:
+            cuda_specs = f"+cuda cuda_arch={arch}"
+            with when(cuda_specs):
+                depends_on(f"py-jaxlib@0.4.13:0.4.26 {cuda_specs}", when="@:0.11.10")
+                depends_on(f"py-jaxlib@0.4.38:0.5 {cuda_specs}", when="@0.11.11:0.11.12")
+                depends_on(f"py-jaxlib@0.4.38:0.6 {cuda_specs}", when="@0.12:")
 
-    with default_args(type=("build", "run")):
-        depends_on("python@3.10:")
         depends_on("py-jax@0.4.13:0.4.26", when="@:0.11.10")
-        depends_on("py-jax@0.4.28:", when="@0.11.11:")
+        depends_on("py-jax@0.4.38:0.5", when="@0.11.11:0.11.12")
+        depends_on("py-jax@0.4.38:0.6", when="@0.12:")
+
         depends_on("py-jaxtyping@0.2.20:")
         depends_on("py-typing-extensions@4.5.0:")
         depends_on("py-wadler-lindig@0.1.0:")
diff --git a/repos/spack_repo/builtin/packages/py_exponax/package.py b/repos/spack_repo/builtin/packages/py_exponax/package.py
@@ -0,0 +1,36 @@
+# Copyright Spack Project Developers. See COPYRIGHT file for details.
+#
+# SPDX-License-Identifier: (Apache-2.0 OR MIT)
+
+
+from spack_repo.builtin.build_systems.cuda import CudaPackage
+from spack_repo.builtin.build_systems.python import PythonPackage
+
+from spack.package import *
+
+
+class PyExponax(PythonPackage, CudaPackage):
+    """Efficient Differentiable n-d PDE solvers built on top of JAX & Equinox."""
+
+    homepage = "https://fkoehler.site/exponax/"
+    pypi = "exponax/exponax-0.1.0.tar.gz"
+
+    maintainers("abhishek1297")
+    license("MIT", checked_by="abhishek1297")
+
+    version("0.1.0", sha256="25acdb5c1b76f5706316750a3133f427f0faec441a1ffe3b90697d5f32abb5e7")
+    version("0.0.1", sha256="e2a201752d38dbfd233d52c2f59ed0dc344ccbb3e796b26c2713c6a2357d7366")
+
+    depends_on("py-setuptools", type="build")
+    depends_on("python@3.10:3.12", type=("build", "run"))
+
+    with default_args(type="run"):
+        for arch in CudaPackage.cuda_arch_values:
+            cuda_specs = f"+cuda cuda_arch={arch}"
+            depends_on(f"py-jaxlib@0.4.13: {cuda_specs}", when=f"{cuda_specs}")
+
+        depends_on("py-jax@0.4.13:")
+        depends_on("py-jaxtyping@0.2.20:")
+        depends_on("py-typing-extensions@4.5.0:")
+        depends_on("py-matplotlib@3.8.1:")
+        depends_on("py-equinox@0.11.3:")
diff --git a/repos/spack_repo/builtin/packages/py_pdequinox/package.py b/repos/spack_repo/builtin/packages/py_pdequinox/package.py
@@ -0,0 +1,37 @@
+# Copyright Spack Project Developers. See COPYRIGHT file for details.
+#
+# SPDX-License-Identifier: (Apache-2.0 OR MIT)
+
+
+from spack_repo.builtin.build_systems.cuda import CudaPackage
+from spack_repo.builtin.build_systems.python import PythonPackage
+
+from spack.package import *
+
+
+class PyPdequinox(PythonPackage, CudaPackage):
+    """A collection of neural architectures for emulating Partial Differential Equations (PDEs)
+    in JAX agnostic to the spatial dimension (1D, 2D, 3D) and boundary conditions
+    (Dirichlet, Neumann, Periodic). This package is built on top of Equinox."""
+
+    homepage = "https://fkoehler.site/pdequinox/"
+    pypi = "pdequinox/pdequinox-0.1.2.tar.gz"
+
+    maintainers("abhishek1297")
+    license("MIT", checked_by="abhishek1297")
+
+    version("0.1.2", sha256="7ee9dcbf277cbb94cda508034c0955600a03bc4c664bede5eb61b4a4b99b54c5")
+    version("0.1.0", sha256="07f7516fe26823e6c3b71f1ed5a170e97cc34ff1d1349435d4b7469adc540d3a")
+
+    depends_on("py-setuptools", type="build")
+    depends_on("python@3.8:3.12", type=("build", "run"))
+
+    with default_args(type="run"):
+        for arch in CudaPackage.cuda_arch_values:
+            cuda_specs = f"+cuda cuda_arch={arch}"
+            depends_on(f"py-jaxlib@0.4.13: {cuda_specs}", when=f"{cuda_specs}")
+
+        depends_on("py-jax@0.4.13:")
+        depends_on("py-jaxtyping@0.2.20:")
+        depends_on("py-typing-extensions@4.5.0:")
+        depends_on("py-equinox@0.11.3:")
diff --git a/repos/spack_repo/builtin/packages/py_trainax/package.py b/repos/spack_repo/builtin/packages/py_trainax/package.py
@@ -0,0 +1,41 @@
+# Copyright Spack Project Developers. See COPYRIGHT file for details.
+#
+# SPDX-License-Identifier: (Apache-2.0 OR MIT)
+
+
+from spack_repo.builtin.build_systems.cuda import CudaPackage
+from spack_repo.builtin.build_systems.python import PythonPackage
+
+from spack.package import *
+
+
+class PyTrainax(PythonPackage, CudaPackage):
+    """Convenience abstractions using optax to train neural networks to autoregressively
+    emulate time-dependent problems taking care of trajectory subsampling and offering a wide
+    range of training methodologies (regarding unrolling length and including
+    differentiable physics).
+    """
+
+    homepage = "https://fkoehler.site/trainax/"
+    pypi = "trainax/trainax-0.0.2.tar.gz"
+
+    maintainers("abhishek1297")
+    license("MIT", checked_by="abhishek1297")
+
+    version("0.0.2", sha256="3c7eeeb94e351db7ff0b036b1c1fb6f78ddc25ab72d6c1afe69547cbefa70ca8")
+    version("0.0.1", sha256="19552dfca2d6f9d7e69963e978628adb19dc2ba9cb9563b510c19e136116c23a")
+
+    depends_on("py-setuptools", type="build")
+    depends_on("python@3.8:3.12", type=("build", "run"))
+
+    with default_args(type="run"):
+        for arch in CudaPackage.cuda_arch_values:
+            cuda_specs = f"+cuda cuda_arch={arch}"
+            depends_on(f"py-jaxlib@0.4.13: {cuda_specs}", when=f"{cuda_specs}")
+
+        depends_on("py-jax@0.4.13:")
+        depends_on("py-jaxtyping@0.2.20:")
+        depends_on("py-typing-extensions@4.5.0:")
+        depends_on("py-tqdm@4.63.2:")
+        depends_on("py-optax@0.2.0:")
+        depends_on("py-equinox@0.11.3:")
