Let TF-GNN support TF2.16+ via TF_USE_LEGACY_KERAS=1.

PiperOrigin-RevId: 619267380

README.md

@@ -52,18 +52,21 @@ page, which also links to the
 The latest stable release of TensorFlow GNN is available from
 
 ```
-pip install tensorflow_gnn
+pip install tensorflow-gnn
 ```
 
 For installation from source, see our [Developer
 Guide](tensorflow_gnn/docs/guide/developer.md).
 
 Key platform requirements:
 
-  * TensorFlow 2.12, 2.13, 2.14 or 2.15, and any GPU drivers it needs
+  * TensorFlow 2.12 or higher, and any GPU drivers it needs
     [[instructions](https://www.tensorflow.org/install)].
   * Keras v2, as traditionally included with TensorFlow 2.x.
-    (TF-GNN does not work with the new multi-backend Keras v3.)
+    TF-GNN does not work with the new multi-backend Keras v3.<br/>
+    **Users of TF2.16+ must also `pip install tf-keras` and set
+    TF_USE_LEGACY_KERAS=1**,
+    see our [Keras version](tensorflow_gnn/docs/guide/keras_version.md) guide for details.
   * Apache Beam for distributed graph sampling.
 
 TF-GNN is developed and tested on Linux. Running on other platforms supported

kokoro/github/ubuntu/cpu/build_versioned.sh

@@ -47,8 +47,11 @@ tag_filters="-no_oss,-oss_excluded${TAG_FILTERS}"
 bazel clean
 pip install -r requirements-dev.txt --progress-bar off
 pip install tensorflow=="${TF_VERSION}" --progress-bar off --upgrade
+if [[ "$TF_USE_LEGACY_KERAS" == 1 ]]; then
+  pip install tf-keras=="${TF_VERSION}" --progress-bar off --upgrade
+fi
 python3 setup.py bdist_wheel
 pip uninstall -y tensorflow_gnn
 pip install dist/tensorflow_gnn-*.whl
 
-bazel test --build_tag_filters="${tag_filters}" --test_tag_filters="${tag_filters}" --test_output=errors --verbose_failures=true --build_tests_only --define=no_tfgnn_py_deps=true --keep_going --experimental_repo_remote_exec //bazel_pip/tensorflow_gnn/...
+bazel test --test_env=TF_USE_LEGACY_KERAS --build_tag_filters="${tag_filters}" --test_tag_filters="${tag_filters}" --test_output=errors --verbose_failures=true --build_tests_only --define=no_tfgnn_py_deps=true --keep_going --experimental_repo_remote_exec //bazel_pip/tensorflow_gnn/...

kokoro/github/ubuntu/cpu/newest_stable/continuous.cfg

@@ -6,7 +6,11 @@ env_vars: {
 }
 env_vars: {
   key: "TF_VERSION"
-  value: "2.15.*"
+  value: "2.16.*"
+}
+env_vars: {
+  key: "TF_USE_LEGACY_KERAS"
+  value: "1"
 }
 
 action {

kokoro/github/ubuntu/cpu/newest_stable/presubmit.cfg

@@ -6,7 +6,11 @@ env_vars: {
 }
 env_vars: {
   key: "TF_VERSION"
-  value: "2.15.*"
+  value: "2.16.*"
+}
+env_vars: {
+  key: "TF_USE_LEGACY_KERAS"
+  value: "1"
 }
 
 action {

kokoro/github/ubuntu/cpu/oldest/continuous.cfg

@@ -8,6 +8,10 @@ env_vars: {
   key: "TF_VERSION"
   value: "2.12.*"
 }
+env_vars: {
+  key: "TF_USE_LEGACY_KERAS"
+  value: "0"
+}
 env_vars: {
   key: "TAG_FILTERS"
   value: ",-tf_at_least_2_13"

kokoro/github/ubuntu/cpu/oldest/presubmit.cfg

@@ -8,6 +8,10 @@ env_vars: {
   key: "TF_VERSION"
   value: "2.12.*"
 }
+env_vars: {
+  key: "TF_USE_LEGACY_KERAS"
+  value: "0"
+}
 env_vars: {
   key: "TAG_FILTERS"
   value: ",-tf_at_least_2_13"

setup.py

@@ -173,8 +173,8 @@ def get_version():
         'networkx',
         'pyarrow',
         # pylint:disable=g-line-too-long
-        'tensorflow>=2.12.0,<2.16.0; platform_machine != "arm64" or platform_system != "Darwin"',
-        'tensorflow-macos>=2.12.0,<2.16.0; platform_machine == "arm64" and platform_system == "Darwin"',
+        'tensorflow>=2.12.0,<3; platform_machine != "arm64" or platform_system != "Darwin"',
+        'tensorflow-macos>=2.12.0,<3; platform_machine == "arm64" and platform_system == "Darwin"',
         # pylint:enable=g-line-too-long
         'apache-beam',
     ],

tensorflow_gnn/__init__.py

@@ -36,7 +36,12 @@ def _check_keras_version():
     if keras_version.startswith("3."):
       raise ImportError(
           "Package tensorflow_gnn requires tf.keras to be Keras version 2 "
-          f"but got version {keras_version}. Install tensorflow<=2.15 to fix.")
+          f"but got version {keras_version}. "
+          "For open-source TensorFlow 2.16 and above, "
+          "set the environment variable TF_USE_LEGACY_KERAS=1 to fix. "
+          "For more information, see "
+          "https://github.com/tensorflow/gnn/blob/main/tensorflow_gnn/docs/guide/keras_version.md"
+      )
 _check_keras_version()
 del _check_keras_version
 # pylint: enable=g-statement-before-imports

tensorflow_gnn/docs/guide/developer.md

@@ -49,6 +49,9 @@ pip install -r requirements-dev.txt
 pip install tensorflow
 ```
 
+For TF2.16+, you will additionally need to follow the instructions from the
+[Keras version](./keras_version.md) guide.
+
 ## Building the package and running tests
 
 We use [Bazel](https://bazel.build/) to build and run tests. Since we use
@@ -84,7 +87,7 @@ After setting up the test directory, you can run all tests locally by running
 the following command in the repo root directory.
 
 ```
-bazel test --build_tag_filters=-no_oss,-oss_excluded --test_tag_filters=-no_oss,-oss_excluded --test_output=errors --verbose_failures=true --build_tests_only --define=no_tfgnn_py_deps=true --keep_going --experimental_repo_remote_exec //bazel_pip/tensorflow_gnn/...
+bazel test --build_tag_filters=-no_oss,-oss_excluded --test_tag_filters=-no_oss,-oss_excluded --test_output=errors --verbose_failures=true --build_tests_only --define=no_tfgnn_py_deps=true --keep_going --experimental_repo_remote_exec --test_env="TF_USE_LEGACY_KERAS=1" //bazel_pip/tensorflow_gnn/...
 ```
 
 The `--define=no_tfgnn_py_deps=true` flag directs bazel to assume that all
@@ -94,19 +97,24 @@ The flags `--build_tag_filters=-no_oss,-oss_excluded` and
 `--test_tag_filters=-no_oss,-oss_excluded` disable tests that pass in the
 internal production environment but fail on GitHub.
 
+The flag `--test_env="TF_USE_LEGACY_KERAS=1"` comes from the
+[Keras version](./keras_version.md) guide and is required for TF2.16+.
+
 ### Run a single test file
 
 It is also possible to run a single test file by specifying its path (for
 example)
 
 `bazel test --define=no_tfgnn_py_deps=true --experimental_repo_remote_exec
+--test_env="TF_USE_LEGACY_KERAS=1"
 //bazel_pip/tensorflow_gnn/models/gcn:gcn_conv_test`
 
 ### Run a single test case
 
 To run a single test case, use the `--test_filter` flag.
 
 `bazel test --define=no_tfgnn_py_deps=true --experimental_repo_remote_exec
+--test_env="TF_USE_LEGACY_KERAS=1"
 //bazel_pip/tensorflow_gnn/models/gcn:gcn_conv_test
 --test_filter=*test_gcnconv_activation*`
 

tensorflow_gnn/docs/guide/keras_version.md

@@ -0,0 +1,87 @@
+# Keras Version Configuration for TF-GNN
+
+The TensorFlow GNN library requires `tf.keras` to be Keras v2, because Keras v3
+does not support composite tensor types like `tfgnn.GraphTensor`. Up to TF 2.15,
+Keras v2 was the default for `tf.keras`. For TF 2.16+, you need to make special
+arrangements, as described in this guide.
+
+<!-- PLACEHOLDER FOR KERAS VERSION GOOGLE EXTRAS -->
+
+## Installation and program execution
+
+### For TensorFlow 2.16 and up
+
+TensorFlow as of release 2.16 depends on `keras>=3` but no longer on a package
+for Keras v2. Install TF-GNN together with the package `tf-keras` that continues
+to supply an implementation of Keras v2.
+
+```
+pip install tensorflow-gnn tf-keras
+```
+
+Running a TF-GNN program under TF 2.16+ requires to set the
+[environment variable](https://en.wikipedia.org/wiki/Environment_variable)
+`TF_USE_LEGACY_KERAS` to `1` one way or another, for example:
+
+  * with the Unix shell command
+
+    ```
+    export TF_USE_LEGACY_KERAS=1
+    ```
+
+    or your system's equivalent;
+
+  * at the top of the main Python program or Colab notebook by
+
+    ```python
+    import os
+    os.environ["TF_USE_LEGACY_KERAS"] = "1"
+    ```
+
+    **before** any part of the program gets to `import tensorflow`.
+
+
+### For TensorFlow 2.15
+
+Nothing special is required: TF 2.15 depends on `keras==2.15.*` and, by default,
+defines `tf.keras` with it.
+
+In case `TF_USE_LEGACY_KERAS` is set to `1` and `tf-keras` is installed, then
+`tf.keras` is defined in terms of that package. While it offers the same API,
+it is a separate package with separate static registries and class hierarchies,
+so it does not mix well with user code that uses `import keras` and objects
+from `keras.*`. Hence we recommend to not use `keras.*` in user code.
+
+### For TensorFlow 2.12, 2.13 and 2.14
+
+Nothing special is required: These TF 2.x versions depend on the matching
+version of Keras 2.x and define `tf.keras` with it. The environment
+variable `TF_USE_LEGACY_KERAS` is ignored.
+
+
+## Writing compatible code
+
+### For all supported TensorFlow versions
+
+For use of TF-GNN under any supported version of TensorFlow, we recommend that
+user code does `import tensorflow as tf` and uses the Keras API at `tf.keras.*`.
+The installation instructions above make sure this is Keras v2, as required
+by the TF-GNN library itself.
+
+Do not import and use `keras.*` directly: it will break for TF 2.16 and above.
+
+### For TF 2.16+ only
+
+As of TF 2.16+, `import tensorflow_gnn as tfgnn` checks that `tf.keras` is
+Keras v2, which implies that it comes from the `tf_keras` package.
+
+Users who wish to emphasize the use of Keras v2 at the expense of breaking
+compatibility with older versions of TensorFlow can use `tf_keras` as a
+syntactic alternative to `tf.keras`.
+
+
+## Other libraries from the TensorFlow ecosystem
+
+Keras does not support mixing different versions or packages of it.
+Using TF-GNN in combination with other TensorFlow add-on libraries requires
+all of them to work with the same Keras package, and it must be Keras v2.

tensorflow_gnn/docs/guide/overview.md

@@ -63,6 +63,9 @@ The following docs go deeper into particular topics.
     saving TF-GNN models. (Most users of TF/Keras 2.13+ should be fine calling
     `tf.keras.Model.export()` without looking here.)
 
+  * The [Keras version config](keras_version.md) guide explains how to install
+    and use Keras v2 with TF2.16 and above, which is required for TF-GNN.
+
 ## Colab Tutorials
 
 These Colab notebooks run complete examples of building and training a TF-GNN