TensorFlow/Keras/TensorBoard: add v2.14.0 (#40297)

Co-authored-by: adamjstewart <adamjstewart@users.noreply.github.com>

4 files changed, 96 insertions, 125 deletions

diff --git a/var/spack/repos/builtin/packages/py-keras/package.py b/var/spack/repos/builtin/packages/py-keras/package.py
index 5604adf859..c6f65dc957 100644
--- a/var/spack/repos/builtin/packages/py-keras/package.py
+++ b/var/spack/repos/builtin/packages/py-keras/package.py
@@ -21,6 +21,7 @@ class PyKeras(PythonPackage):
     git = "https://github.com/keras-team/keras.git"
     url = "https://github.com/keras-team/keras/archive/refs/tags/v2.7.0.tar.gz"
 
+    version("2.14.0", sha256="a845d446b6ae626f61dde5ab2fa952530b6c17b4f9ed03e9362bd20172d00cca")
     version("2.13.1", sha256="b3591493cce75a69adef7b192cec6be222e76e2386d132cd4e34aa190b0ecbd5")
     version("2.12.0", sha256="6336cebb6b2b0a91f7efd3ff3a9db3a94f2abccf07a40323138afb80826aec62")
     version("2.11.0", sha256="e7a7c4199ac76ea750d145c1d84ae1b932e68b9bca34e83596bd66b2fc2ad79e")
@@ -61,7 +62,7 @@ class PyKeras(PythonPackage):
     depends_on("py-pydot", type=("build", "run"))
     depends_on("py-scipy", type=("build", "run"))
     depends_on("py-six", type=("build", "run"))
-    for minor_ver in range(6, 14):
+    for minor_ver in range(6, 15):
         depends_on(
             "py-tensorflow@2.{}".format(minor_ver),
             type=("build", "run"),
diff --git a/var/spack/repos/builtin/packages/py-ml-dtypes/package.py b/var/spack/repos/builtin/packages/py-ml-dtypes/package.py
new file mode 100644
index 0000000000..192069e219
--- /dev/null
+++ b/var/spack/repos/builtin/packages/py-ml-dtypes/package.py
@@ -0,0 +1,25 @@
+# Copyright 2013-2023 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 PyMlDtypes(PythonPackage):
+    """A stand-alone implementation of several NumPy dtype extensions
+    used in machine learning libraries."""
+
+    homepage = "https://github.com/jax-ml/ml_dtypes"
+    pypi = "ml_dtypes/ml_dtypes-0.3.1.tar.gz"
+    git = "https://github.com/jax-ml/ml_dtypes.git"
+    submodules = True
+
+    version("0.3.1", tag="v0.3.1", commit="bbeedd470ecac727c42e97648c0f27bfc312af30")
+    version("0.2.0", tag="v0.2.0", commit="5b9fc9ad978757654843f4a8d899715dbea30e88")
+
+    depends_on("python@3.9:", when="@0.3:", type=("build", "link", "run"))
+    depends_on("py-numpy@1.21:", type=("build", "link", "run"))
+    # Build dependencies are overconstrained, older versions work just fine
+    depends_on("py-pybind11", type=("build", "link"))
+    depends_on("py-setuptools", type="build")
diff --git a/var/spack/repos/builtin/packages/py-tensorboard/package.py b/var/spack/repos/builtin/packages/py-tensorboard/package.py
index f1cd059f69..f1a9b03bc1 100644
--- a/var/spack/repos/builtin/packages/py-tensorboard/package.py
+++ b/var/spack/repos/builtin/packages/py-tensorboard/package.py
@@ -7,9 +7,8 @@ from spack.package import *
 
 
 class PyTensorboard(PythonPackage):
-    """TensorBoard is a suite of web applications for
-    inspecting and understanding your TensorFlow runs and
-    graphs."""
+    """TensorBoard is a suite of web applications for inspecting and understanding
+    your TensorFlow runs and graphs."""
 
     homepage = "https://github.com/tensorflow/tensorboard"
     url = "https://files.pythonhosted.org/packages/py3/t/tensorboard/tensorboard-2.9.1-py3-none-any.whl"
@@ -18,6 +17,16 @@ class PyTensorboard(PythonPackage):
     maintainers("aweits")
 
     version(
+        "2.14.1",
+        sha256="3db108fb58f023b6439880e177743c5f1e703e9eeb5fb7d597871f949f85fd58",
+        expand=False,
+    )
+    version(
+        "2.14.0",
+        sha256="3667f9745d99280836ad673022362c840f60ed8fefd5a3e30bf071f5a8fd0017",
+        expand=False,
+    )
+    version(
         "2.13.0",
         sha256="ab69961ebddbddc83f5fa2ff9233572bdad5b883778c35e4fe94bf1798bd8481",
         expand=False,
@@ -118,6 +127,7 @@ class PyTensorboard(PythonPackage):
         expand=False,
     )
 
+    depends_on("python@3.9:", type=("build", "run"), when="@2.14:")
     depends_on("python@3.8:", type=("build", "run"), when="@2.12:")
     depends_on("py-absl-py@0.4:", type=("build", "run"))
     depends_on("py-grpcio@1.48.2:", type=("build", "run"), when="@2.12:")
@@ -135,10 +145,10 @@ class PyTensorboard(PythonPackage):
     depends_on("py-protobuf@3.6.0:3.19", type=("build", "run"), when="@:2.8")
     depends_on("py-requests@2.21.0:2", type=("build", "run"))
     depends_on("py-setuptools@41.0.0:", type=("build", "run"))
+    depends_on("py-six@1.10.0:", type=("build", "run"), when="@:2.4,2.14:")
     depends_on("py-tensorboard-data-server@0.7", type=("build", "run"), when="@2.12:")
     depends_on("py-tensorboard-data-server@0.6", type=("build", "run"), when="@2.5:2.11")
-    depends_on("py-tensorboard-plugin-wit@1.6.0:", type=("build", "run"))
+    depends_on("py-tensorboard-plugin-wit@1.6.0:", type=("build", "run"), when="@:2.13")
     depends_on("py-werkzeug@1.0.1:", type=("build", "run"), when="@2.9:")
     depends_on("py-werkzeug@0.11.15:", type=("build", "run"))
-    depends_on("py-wheel@0.26:", type="build")
-    depends_on("py-six@1.10.0:", type=("build", "run"), when="@:2.4")
+    depends_on("py-wheel@0.26:", type="build", when="@:2.13")
diff --git a/var/spack/repos/builtin/packages/py-tensorflow/package.py b/var/spack/repos/builtin/packages/py-tensorflow/package.py
index cd90b532c9..c85f078689 100644
--- a/var/spack/repos/builtin/packages/py-tensorflow/package.py
+++ b/var/spack/repos/builtin/packages/py-tensorflow/package.py
@@ -10,18 +10,7 @@ from spack.package import *
 
 
 class PyTensorflow(Package, CudaPackage, ROCmPackage, PythonExtension):
-    """An Open Source Machine Learning Framework for Everyone.
-
-    TensorFlow is an end-to-end open source platform for machine learning. It has a
-    comprehensive, flexible ecosystem of tools, libraries, and community resources that
-    lets researchers push the state-of-the-art in ML and developers easily build and
-    deploy ML-powered applications.
-
-    TensorFlow was originally developed by researchers and engineers working on the
-    Google Brain team within Google's Machine Intelligence Research organization to
-    conduct machine learning and deep neural networks research. The system is general
-    enough to be applicable in a wide variety of other domains, as well.
-    """
+    """TensorFlow is an open source machine learning framework for everyone."""
 
     homepage = "https://www.tensorflow.org"
     url = "https://github.com/tensorflow/tensorflow/archive/v2.3.1.tar.gz"
@@ -29,6 +18,8 @@ class PyTensorflow(Package, CudaPackage, ROCmPackage, PythonExtension):
     maintainers("adamjstewart", "aweits")
     import_modules = ["tensorflow"]
 
+    version("2.14.0", sha256="ce357fd0728f0d1b0831d1653f475591662ec5bca736a94ff789e6b1944df19f")
+    version("2.13.1", sha256="89c07aebd4f41fbe0d08cc88aef00305542134f2f16d3b62918dc3c1182f33e2")
     version("2.13.0", sha256="e58c939079588623e6fa1d054aec2f90f95018266e0a970fd353a5244f5173dc")
     version("2.12.1", sha256="6bc4600cc0b88e9e40f1800096f5bddbbd3b6e5527a030dea631b87f2ae46b5b")
     version("2.12.0", sha256="c030cb1905bff1d2446615992aad8d8d85cbe90c4fb625cee458c63bf466bc8e")
@@ -157,16 +148,16 @@ class PyTensorflow(Package, CudaPackage, ROCmPackage, PythonExtension):
     extends("python")
 
     # Python support based on wheel availability
-    depends_on("python@3.8:3.11", when="@2.12:", type=("build", "run"))
-    depends_on("python@3.7:3.10", when="@2.8:2.11", type=("build", "run"))
-    depends_on("python@3.7:3.9", when="@2.7", type=("build", "run"))
-    depends_on("python@3.6:3.9", when="@2.5:2.6", type=("build", "run"))
-    depends_on("python@3.6:3.8", when="@2.4", type=("build", "run"))
-    depends_on("python@3.5:3.8", when="@2.2:2.3", type=("build", "run"))
-    depends_on("python@2.7,3.5:3.7", when="@:2.1", type=("build", "run"))
+    depends_on("python@3.9:3.11", when="@2.14:", type=("build", "run"))
+    depends_on("python@3.8:3.11", when="@2.12:2.13", type=("build", "run"))
+    depends_on("python@:3.10", when="@2.8:2.11", type=("build", "run"))
+    depends_on("python@:3.9", when="@2.5:2.7", type=("build", "run"))
+    depends_on("python@:3.8", when="@2.2:2.4", type=("build", "run"))
+    depends_on("python@:3.7", when="@:2.1", type=("build", "run"))
 
     # See .bazelversion
-    depends_on("bazel@5.3.0", type="build", when="@2.11:")
+    depends_on("bazel@6.1.0", type="build", when="@2.14:")
+    depends_on("bazel@5.3.0", type="build", when="@2.11:2.13")
     depends_on("bazel@5.1.1", type="build", when="@2.10")
     # See _TF_MIN_BAZEL_VERSION and _TF_MAX_BAZEL_VERSION in configure.py
     depends_on("bazel@4.2.2:5.99.0", type="build", when="@2.9")
@@ -203,13 +194,15 @@ class PyTensorflow(Package, CudaPackage, ROCmPackage, PythonExtension):
     depends_on("py-astunparse@1.6:", type=("build", "run"), when="@2.7:")
     depends_on("py-astunparse@1.6.3:1.6", type=("build", "run"), when="@2.4:2.6")
     depends_on("py-astunparse@1.6.3", type=("build", "run"), when="@2.2:2.3")
-    depends_on("py-flatbuffers@23.1.21:", type=("build", "run"), when="@2.13:")
+    depends_on("py-flatbuffers@23.5.26:", type=("build", "run"), when="@2.14:")
+    depends_on("py-flatbuffers@23.1.21:", type=("build", "run"), when="@2.13")
     depends_on("py-flatbuffers@2:", type=("build", "run"), when="@2.10:2.12")
     depends_on("py-flatbuffers@1.12:1", type=("build", "run"), when="@2.9")
     depends_on("py-flatbuffers@1.12:", type=("build", "run"), when="@2.8")
     depends_on("py-flatbuffers@1.12:2", type=("build", "run"), when="@2.7")
     depends_on("py-flatbuffers@1.12", type=("build", "run"), when="@2.4:2.6")
-    depends_on("py-gast@0.2.1:0.4.0", type=("build", "run"), when="@2.9:")
+    depends_on("py-gast@0.2.1:0.4,0.5.3:", type=("build", "run"), when="@2.14:")
+    depends_on("py-gast@0.2.1:0.4.0", type=("build", "run"), when="@2.9:2.13")
     depends_on("py-gast@0.2.1:", type=("build", "run"), when="@2.8")
     depends_on("py-gast@0.2.1:0.4", type=("build", "run"), when="@2.7")
     depends_on("py-gast@0.4.0", type=("build", "run"), when="@2.5:2.6")
@@ -232,6 +225,8 @@ class PyTensorflow(Package, CudaPackage, ROCmPackage, PythonExtension):
     depends_on("hdf5~mpi", type="build", when="@1.15.5,2.0.4,2.1.3:~mpi")
     depends_on("py-libclang@13:", type=("build", "run"), when="@2.9:")
     depends_on("py-libclang@9.0.1:", type=("build", "run"), when="@2.7:2.8")
+    depends_on("py-ml-dtypes@0.2.0", type=("build", "run"), when="@2.14:")
+    depends_on("py-numpy@1.23.5:", type=("build", "run"), when="@2.14:")
     depends_on("py-numpy@1.22:1.24.3", type=("build", "run"), when="@2.13:")
     depends_on("py-numpy@1.22:1.23", type=("build", "run"), when="@2.12")
     depends_on("py-numpy@1.20:", type=("build", "run"), when="@2.8:2.11")
@@ -269,32 +264,17 @@ class PyTensorflow(Package, CudaPackage, ROCmPackage, PythonExtension):
     depends_on("py-six@1.10:", type=("build", "run"), when="@:2.0")
     depends_on("py-termcolor@1.1:", type=("build", "run"), when="@1.6:2.3,2.7:")
     depends_on("py-termcolor@1.1", type=("build", "run"), when="@2.4:2.6")
-    depends_on("py-typing-extensions@3.6.6:4.5", type=("build", "run"), when="@2.13:")
-    depends_on("py-typing-extensions@3.6.6:", type=("build", "run"), when="@2.7:2.12")
+    depends_on("py-typing-extensions@3.6.6:", type=("build", "run"), when="@2.7:2.12,2.14:")
+    depends_on("py-typing-extensions@3.6.6:4.5", type=("build", "run"), when="@2.13")
     depends_on("py-typing-extensions@3.7.4:3.7", type=("build", "run"), when="@2.4:2.6")
-    depends_on("py-wrapt@1.11:", type=("build", "run"), when="@2.13:")
-    depends_on("py-wrapt@1.11:1.14", type=("build", "run"), when="@2.12")
-    depends_on("py-wrapt@1.11:", type=("build", "run"), when="@2.7:2.11")
+    depends_on("py-wrapt@1.11:1.14", type=("build", "run"), when="@2.12,2.14:")
+    depends_on("py-wrapt@1.11:", type=("build", "run"), when="@2.7:2.11,2.13")
     depends_on("py-wrapt@1.12.1:1.12", type=("build", "run"), when="@2.4:2.6")
     depends_on("py-wrapt@1.11.1:", type=("build", "run"), when="@1.12.1,1.14:2.3")
+
     # TODO: add packages for these dependencies
     # depends_on('py-tensorflow-io-gcs-filesystem@0.23.1:', type=('build', 'run'), when='@2.8:')
     # depends_on('py-tensorflow-io-gcs-filesystem@0.21:', type=('build', 'run'), when='@2.7')
-    with when("+rocm"):
-        depends_on("hip")
-        depends_on("rocrand")
-        depends_on("rocblas")
-        depends_on("rocfft")
-        depends_on("hipfft")
-        depends_on("rccl", when="+nccl")
-        depends_on("hipsparse")
-        depends_on("hipcub")
-        depends_on("rocsolver")
-        depends_on("rocprim")
-        depends_on("miopen-hip")
-        depends_on("llvm-amdgpu")
-        depends_on("hsa-rocr-dev")
-        depends_on("rocminfo")
 
     if sys.byteorder == "little":
         # Only builds correctly on little-endian machines
@@ -304,7 +284,7 @@ class PyTensorflow(Package, CudaPackage, ROCmPackage, PythonExtension):
         depends_on("py-grpcio@1.32", type=("build", "run"), when="@2.4")
         depends_on("py-grpcio@1.8.6:", type=("build", "run"), when="@1.6:2.3")
 
-    for minor_ver in range(5, 14):
+    for minor_ver in range(5, 15):
         depends_on(
             "py-tensorboard@2.{}".format(minor_ver),
             type=("build", "run"),
@@ -356,6 +336,22 @@ class PyTensorflow(Package, CudaPackage, ROCmPackage, PythonExtension):
     # depends_on('android-ndk@10:18', when='+android')
     # depends_on('android-sdk', when='+android')
 
+    with when("+rocm"):
+        depends_on("hip")
+        depends_on("rocrand")
+        depends_on("rocblas")
+        depends_on("rocfft")
+        depends_on("hipfft")
+        depends_on("rccl", when="+nccl")
+        depends_on("hipsparse")
+        depends_on("hipcub")
+        depends_on("rocsolver")
+        depends_on("rocprim")
+        depends_on("miopen-hip")
+        depends_on("llvm-amdgpu")
+        depends_on("hsa-rocr-dev")
+        depends_on("rocminfo")
+
     # Check configure and configure.py to see when these variants are supported
     conflicts("+mkl", when="@:1.0")
     conflicts("+mkl", when="platform=darwin", msg="Darwin is not yet supported")
@@ -375,7 +371,13 @@ class PyTensorflow(Package, CudaPackage, ROCmPackage, PythonExtension):
     conflicts("+gdr", when="@:1.3")
     conflicts("+verbs", when="@:1.1")
     conflicts("+ngraph", when="@:1.10")
+    conflicts("+opencl", when="platform=windows")
     conflicts("+computecpp", when="~opencl")
+    conflicts(
+        "+cuda",
+        when="+rocm",
+        msg="CUDA / ROCm are mututally exclusive. At most 1 GPU platform can be configured",
+    )
     conflicts("+cuda", when="platform=darwin", msg="There is no GPU support for macOS")
     conflicts(
         "cuda_arch=none",
@@ -422,6 +424,7 @@ class PyTensorflow(Package, CudaPackage, ROCmPackage, PythonExtension):
     conflicts(
         "+nccl", when="platform=cray", msg="Currently NCCL is only supported on Linux platform"
     )
+    conflicts("+mpi", when="platform=windows")
     conflicts("+mpi", when="@:1.2")
     conflicts("+android", when="@:1.4")
     conflicts("+ios", when="@:1.12.0,1.12.2:1.13")
@@ -515,6 +518,7 @@ class PyTensorflow(Package, CudaPackage, ROCmPackage, PythonExtension):
 
         # Please input the desired Python library path to use
         env.set("PYTHON_LIB_PATH", python_platlib)
+        env.set("TF_PYTHON_VERSION", spec["python"].version.up_to(2))
 
         # Ensure swig is in PATH or set SWIG_PATH
         env.set("SWIG_PATH", spec["swig"].prefix.bin.swig)
@@ -696,6 +700,12 @@ class PyTensorflow(Package, CudaPackage, ROCmPackage, PythonExtension):
         else:
             env.set("TF_NEED_CUDA", "0")
 
+        # Do you want to use Clang to build TensorFlow?
+        if "%clang" in spec:
+            env.set("TF_NEED_CLANG", "1")
+        else:
+            env.set("TF_NEED_CLANG", "0")
+
         # Do you wish to download a fresh release of clang? (Experimental)
         env.set("TF_DOWNLOAD_CLANG", "0")
 
@@ -789,14 +799,6 @@ class PyTensorflow(Package, CudaPackage, ROCmPackage, PythonExtension):
                 "tensorflow/workspace.bzl",
             )
 
-            # starting with tensorflow 1.3, tensorboard becomes a dependency
-            # -> remove from list of required packages
-            filter_file(
-                r"'tensorflow-tensorboard",
-                r"#'tensorflow-tensorboard",
-                "tensorflow/tools/pip_package/setup.py",
-            )
-
         if spec.satisfies("@1.5.0: ~gcp"):
             # google cloud support seems to be installed on default, leading
             # to boringssl error manually set the flag to false to avoid
@@ -808,15 +810,6 @@ class PyTensorflow(Package, CudaPackage, ROCmPackage, PythonExtension):
                 ".tf_configure.bazelrc",
             )
 
-        if spec.satisfies("@1.6.0:2.1"):
-            # tensorboard name changed
-            # there are no corresponding versions of these in spack
-            filter_file(
-                r"(^\s*)'tensorboard (>=|~=)",
-                r"\1#'tensorboard \2",
-                "tensorflow/tools/pip_package/setup.py",
-            )
-
         if spec.satisfies("@1.8.0: ~opencl"):
             # 1.8.0 and 1.9.0 aborts with numpy import error during python_api
             # generation somehow the wrong PYTHONPATH is used...
@@ -826,64 +819,6 @@ class PyTensorflow(Package, CudaPackage, ROCmPackage, PythonExtension):
                 f.write("build --distinct_host_configuration=false\n")
                 f.write('build --action_env PYTHONPATH="{0}"\n'.format(env["PYTHONPATH"]))
 
-        if spec.satisfies("@1.13.1:"):
-            # tensorflow_estimator is an API for tensorflow
-            # tensorflow-estimator imports tensorflow during build, so
-            # tensorflow has to be set up first
-            filter_file(
-                r"(^\s*)'tensorflow_estimator (>=|~=)",
-                r"\1#'tensorflow_estimator \2",
-                "tensorflow/tools/pip_package/setup.py",
-            )
-
-        if spec.satisfies("@2.5"):
-            filter_file(
-                r"(^\s*)'keras-nightly (>=|~=)",
-                r"\1#'keras-nightly \2",
-                "tensorflow/tools/pip_package/setup.py",
-            )
-
-        if spec.satisfies("@2.6:"):
-            filter_file(
-                r"(^\s*)'keras (>=|~=)", r"\1#'keras \2", "tensorflow/tools/pip_package/setup.py"
-            )
-
-        if spec.satisfies("@2.6"):
-            filter_file(
-                r"(^\s*)'clang (>=|~=)", r"\1#'clang \2", "tensorflow/tools/pip_package/setup.py"
-            )
-
-        # TODO: add support for tensorflow-io-gcs-filesystem
-        if spec.satisfies("@2.7:"):
-            filter_file(
-                r"(^\s*)'tensorflow-io-gcs-filesystem (>=|~=)",
-                r"\1#'tensorflow-io-gcs-filesystem \2",
-                "tensorflow/tools/pip_package/setup.py",
-            )
-
-        if spec.satisfies("@2.0.0:"):
-            # now it depends on the nightly versions...
-            filter_file(
-                r"REQUIRED_PACKAGES\[i\] = 'tb-nightly (>=|~=)",
-                r"pass #REQUIRED_PACKAGES[i] = 'tb-nightly \1",
-                "tensorflow/tools/pip_package/setup.py",
-            )
-            filter_file(
-                r"REQUIRED_PACKAGES\[i\] = 'tensorflow-estimator-2.0-preview",
-                r"pass #REQUIRED_PACKAGES[i] = 'tensorflow-estimator-2.0-preview",
-                "tensorflow/tools/pip_package/setup.py",
-            )
-            filter_file(
-                r"REQUIRED_PACKAGES\[i\] = 'tf-estimator-nightly (>=|~=)",
-                r"pass #REQUIRED_PACKAGES[i] = 'tf-estimator-nightly \1",
-                "tensorflow/tools/pip_package/setup.py",
-            )
-            filter_file(
-                r"REQUIRED_PACKAGES\[i\] = 'keras-nightly (>=|~=)",
-                r"pass #REQUIRED_PACKAGES[i] = 'keras-nightly \1",
-                "tensorflow/tools/pip_package/setup.py",
-            )
-
         if spec.satisfies("@1.13.1 +nccl"):
             filter_file(
                 r"^build --action_env NCCL_INSTALL_PATH=.*",