Add py-tf-keras package, upgrade TFP (#43688)

* enh: add tf-keras package, upgrade TFP

* chore:  remove legacy deps

* chore:  fix style

* chore:  fix style

* fix: url

* fix: use jax, tensorflow instead of py-jax, py-tensorflow

* fix: remove typo

* Update var/spack/repos/builtin/packages/py-tensorflow-probability/package.py

Co-authored-by: Adam J. Stewart <ajstewart426@gmail.com>

* fix: typos

* fix: swap version

* fix: typos

* fix: typos

* fix: typos

* chore: use f strings

* enh: move tf-keras to pypi

* [@spackbot] updating style on behalf of jonas-eschle

* fix: t

* enh: add tf-keras package, upgrade TFP

* chore:  remove legacy deps

* chore:  fix style

* chore:  fix style

* fix: url

* fix: use jax, tensorflow instead of py-jax, py-tensorflow

* fix: remove typo

* Update var/spack/repos/builtin/packages/py-tensorflow-probability/package.py

Co-authored-by: Adam J. Stewart <ajstewart426@gmail.com>

* fix: typos

* fix: swap version

* fix: typos

* fix: typos

* fix: typos

* chore: use f strings

* enh: move tf-keras to pypi

* [@spackbot] updating style on behalf of jonas-eschle

* enh: move tf-keras to pypi

* enh: move back to releases to make it work, actually

* enh: move back to releases to make it work, actually

* fix:change back to tar...

* Fix concretisation: py-tf-keras only has 2.17, not 2.16, fix checksum

* enh: add TFP 0.25

* enh: add tf-keras 2.18

* chore: fix style

* fix: remove patch

* maybe fix license

* Update var/spack/repos/builtin/packages/py-tf-keras/package.py

Co-authored-by: Adam J. Stewart <ajstewart426@gmail.com>

* fix:  pipargs global?

* Update var/spack/repos/builtin/packages/py-tf-keras/package.py

Co-authored-by: Wouter Deconinck <wdconinc@gmail.com>

* chore: fix formatting

* chore: fix formatting again

* fix:  pathes in spack

* fix:  typo

* fix:  typo

* use github package

* use pip install

* fix typo

* fix typo

* comment 2.19 out

* fix typo

* fix typo

* fix typo

* chore: remove unused patch file

* chore: cleanup

* chore: add comment about TF version

* chore: remove unused Bazel, cleanup imports

* [@spackbot] updating style on behalf of jonas-eschle

* chore: add star import, degrading readability

---------

Co-authored-by: Adam J. Stewart <ajstewart426@gmail.com>
Co-authored-by: jonas-eschle <jonas-eschle@users.noreply.github.com>
Co-authored-by: Bernhard Kaindl <contact@bernhard.kaindl.dev>
Co-authored-by: Bernhard Kaindl <bernhardkaindl7@gmail.com>
Co-authored-by: Wouter Deconinck <wdconinc@gmail.com>

var/spack/repos/builtin/packages/py-tensorflow-probability/package.py

@@ -17,14 +17,14 @@ class PyTensorflowProbability(Package):
     predictions."""
 
     homepage = "https://www.tensorflow.org/probability"
-    url = "https://github.com/tensorflow/probability/archive/v0.12.1.tar.gz"
+    url = "https://github.com/tensorflow/probability/archive/v0.24.0.tar.gz"
 
     maintainers("aweits", "jonas-eschle")
 
     license("Apache-2.0")
 
-    # tf-keras does not (yet) exist on Spack, i.e. the TF extra won't work, see below
+    version("0.25.0", sha256="73cf3e52d3b5ee48fc919631e7fa3621849b8ce8d453017016a5148815856bc2")
-    # version("0.24.0", sha256="3d75418cef09ea357ee879347133ab784fe4637a5ba251a8e06ef930dd970a3e")
+    version("0.24.0", sha256="3d75418cef09ea357ee879347133ab784fe4637a5ba251a8e06ef930dd970a3e")
     version("0.23.0", sha256="a00769550da9284acbd69e32a005507153ad39b0c190feca2bbbf6373366cc14")
     version("0.22.1", sha256="9c1203b454aaeb48ac67dea862a411dba6b04f67c1e874e0e83bd1d7f13829a3")
     version("0.22.0", sha256="f9ce55b00c8069246d701c04eaafccde413355f6e76ccf9e549772ecfa0349a4")
@@ -58,14 +58,11 @@ class PyTensorflowProbability(Package):
     depends_on("py-dm-tree", when="@0.12:", type=("build", "run"))
 
     # tensorflow_probability/python/__init__.py
-    # TODO: reactivate the JAX versions once the JAX package is available with newer versions
+    variant("tf", default=False, description="Build with TensorFlow support")
-    #  also add jaxlib as a dependency
+    with when("+tf"):
-
+        depends_on("py-tf-keras@2.18:", when="@0.25:", type=("build", "run"))
-    variant("py-tensorflow", default=False, description="Build with TensorFlow support")
+        depends_on("py-tf-keras@2.17:", when="@0.24:", type=("build", "run"))
-    with when("+py-tensorflow"):
+        depends_on("py-tensorflow@2.17", when="@0.24", type=("build", "run"))
-        # TODO(0.24): requires tf-keras on Spack
-        # depends_on("py-tf-keras", when="@0.24:", type=("build", "run"))
-        # depends_on("py-tensorflow@2.16", when="@0.24", type=("build", "run"))
         depends_on("py-tensorflow@2.15", when="@0.23", type=("build", "run"))
         depends_on("py-tensorflow@2.14:2", when="@0.22", type=("build", "run"))
         depends_on("py-tensorflow@2.13:2", when="@0.21", type=("build", "run"))
@@ -73,8 +70,13 @@ class PyTensorflowProbability(Package):
         depends_on("py-tensorflow@2.11:2", when="@0.19", type=("build", "run"))
 
     # jaxlib is not required, as it's already a dependency of py-jax
-    variant("py-jax", default=False, description="Build with JAX support")
+    # no real minimum version, setup.py doesn't specify one.
-    with when("+py-jax"):
+    # The README only specifies that it is "tested and stable against".
+    # If needed, this versions could be relaxed.
+    variant("jax", default=False, description="Build with JAX support")
+    with when("+jax"):
+        depends_on("py-jax@0.4.35:0.5", when="@0.25", type=("build", "run"))
+        depends_on("py-jax@0.4.25:0.4", when="@0.24", type=("build", "run"))
         depends_on("py-jax@0.4.20:0.4", when="@0.23", type=("build", "run"))
         depends_on("py-jax@0.4.16:0.4", when="@0.22", type=("build", "run"))
         depends_on("py-jax@0.4.14:0.4", when="@0.21", type=("build", "run"))

var/spack/repos/builtin/packages/py-tf-keras/package.py

@@ -0,0 +1,58 @@
+# Copyright Spack Project Developers. See COPYRIGHT file for details.
+#
+# SPDX-License-Identifier: (Apache-2.0 OR MIT)
+
+from spack.package import *
+
+
+class PyTfKeras(PythonPackage):
+    """The TF-Keras library is a pure TensorFlow implementation of Keras,
+    based on the legacy tf.keras codebase. Note that the "main" version
+    of Keras is now Keras 3 (formerly Keras Core), which is a
+    multi-backend implementation of Keras, supporting JAX, PyTorch, and TensorFlow.
+     Keras 3 is being developed at keras-team/keras."""
+
+    homepage = "https://github.com/keras-team/tf-keras"
+    pypi = "tf-keras/tf_keras-2.18.0.tar.gz"
+    # url = "https://github.com/keras-team/tf-keras/archive/refs/tags/v2.18.0.tar.gz"
+
+    maintainers("jonas-eschle")
+
+    license("Apache-2.0", checked_by="jonas-eschle")
+
+    max_minor = 18
+    # needs TF 2.19: https://github.com/spack/spack/pull/49440
+    # version("2.19.0", sha256="b09a407d87a4571ce1e8ca985cfc68483e3d63b2518a5d79a97ad92cb64dbe9c")
+    version("2.18.0", sha256="ebf744519b322afead33086a2aba872245473294affd40973694f3eb7c7ad77d")
+
+    # Supported Python versions listed in multiple places:
+    # * tf-keras/tools/pip_package/setup.py
+    # * CONTRIBUTING.md
+    # * PKG-INFO
+    depends_on("python@3.9:", type=("build", "run"), when="@2.17:")
+    depends_on("py-setuptools", type="build")
+
+    # Required dependencies listed in multiple places:
+    # * BUILD
+    # * WORKSPACE
+    depends_on("py-pandas", type=("build", "run"))
+    depends_on("py-pydot", type=("build", "run"))
+    depends_on("py-scipy", type=("build", "run"))
+
+    depends_on("protobuf@3.20.3", type="build", when="@2.18:")
+    # TODO: uncomment for 2.19
+    # depends_on("protobuf@4.23.0", type="build", when="@2.19:")
+    # the tf-keras versions are following along with TF versions
+    # as defined in oss_setup.py
+    for minor_ver in range(18, max_minor + 1):
+        depends_on(f"py-tensorflow@2.{minor_ver}", type=("build", "run"), when=f"@2.{minor_ver}")
+        # depends_on(f"py-tensorboard@2.{minor_ver}",
+        # type=("build", "run"), when=f"@2.{minor_ver}")
+    depends_on("py-portpicker", type=("build", "run"))
+    depends_on("py-pyyaml", type=("build", "run"))
+    depends_on("pil", type=("build", "run"))
+    depends_on("py-numpy@1.26.0:2.0", type=("build", "run"), when="@2.18:")
+
+    depends_on("py-six", type=("build", "run"))
+    depends_on("py-absl-py", type=("build", "run"))
+    depends_on("py-h5py", type=("build", "run"))