From c2be3427599dae5f87fc9b022968c26a97a87ef9 Mon Sep 17 00:00:00 2001 From: Arno Eigenwillig Date: Mon, 25 Mar 2024 09:23:02 -0700 Subject: [PATCH] Let TF-GNN support TF2.16+ via `TF_USE_LEGACY_KERAS=1`. PiperOrigin-RevId: 618864989 --- README.md | 9 +- kokoro/github/ubuntu/cpu/build_versioned.sh | 2 +- .../ubuntu/cpu/newest_stable/continuous.cfg | 6 +- .../ubuntu/cpu/newest_stable/presubmit.cfg | 6 +- .../github/ubuntu/cpu/oldest/continuous.cfg | 4 + kokoro/github/ubuntu/cpu/oldest/presubmit.cfg | 4 + setup.py | 4 +- tensorflow_gnn/__init__.py | 7 +- tensorflow_gnn/docs/guide/developer.md | 10 ++- tensorflow_gnn/docs/guide/keras_version.md | 87 +++++++++++++++++++ tensorflow_gnn/docs/guide/overview.md | 3 + 11 files changed, 132 insertions(+), 10 deletions(-) create mode 100644 tensorflow_gnn/docs/guide/keras_version.md diff --git a/README.md b/README.md index 2f53aedc..0e54c250 100644 --- a/README.md +++ b/README.md @@ -52,7 +52,7 @@ 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 @@ -60,10 +60,13 @@ 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.
+ **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 diff --git a/kokoro/github/ubuntu/cpu/build_versioned.sh b/kokoro/github/ubuntu/cpu/build_versioned.sh index 283fee92..2caa6614 100644 --- a/kokoro/github/ubuntu/cpu/build_versioned.sh +++ b/kokoro/github/ubuntu/cpu/build_versioned.sh @@ -51,4 +51,4 @@ 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="${TEST_ENV}" --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/... diff --git a/kokoro/github/ubuntu/cpu/newest_stable/continuous.cfg b/kokoro/github/ubuntu/cpu/newest_stable/continuous.cfg index 12819599..c33bfd2a 100644 --- a/kokoro/github/ubuntu/cpu/newest_stable/continuous.cfg +++ b/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: "TEST_ENV" + value: "TF_USE_LEGACY_KERAS=1" } action { diff --git a/kokoro/github/ubuntu/cpu/newest_stable/presubmit.cfg b/kokoro/github/ubuntu/cpu/newest_stable/presubmit.cfg index 12819599..c33bfd2a 100644 --- a/kokoro/github/ubuntu/cpu/newest_stable/presubmit.cfg +++ b/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: "TEST_ENV" + value: "TF_USE_LEGACY_KERAS=1" } action { diff --git a/kokoro/github/ubuntu/cpu/oldest/continuous.cfg b/kokoro/github/ubuntu/cpu/oldest/continuous.cfg index 1beddb6c..e58917a5 100644 --- a/kokoro/github/ubuntu/cpu/oldest/continuous.cfg +++ b/kokoro/github/ubuntu/cpu/oldest/continuous.cfg @@ -8,6 +8,10 @@ env_vars: { key: "TF_VERSION" value: "2.12.*" } +env_vars: { + key: "TEST_ENV" + value: "" +} env_vars: { key: "TAG_FILTERS" value: ",-tf_at_least_2_13" diff --git a/kokoro/github/ubuntu/cpu/oldest/presubmit.cfg b/kokoro/github/ubuntu/cpu/oldest/presubmit.cfg index 1beddb6c..e58917a5 100644 --- a/kokoro/github/ubuntu/cpu/oldest/presubmit.cfg +++ b/kokoro/github/ubuntu/cpu/oldest/presubmit.cfg @@ -8,6 +8,10 @@ env_vars: { key: "TF_VERSION" value: "2.12.*" } +env_vars: { + key: "TEST_ENV" + value: "" +} env_vars: { key: "TAG_FILTERS" value: ",-tf_at_least_2_13" diff --git a/setup.py b/setup.py index 3beaeef3..9e1144c5 100644 --- a/setup.py +++ b/setup.py @@ -175,8 +175,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', ], diff --git a/tensorflow_gnn/__init__.py b/tensorflow_gnn/__init__.py index 5bab00c7..7a4b2a70 100644 --- a/tensorflow_gnn/__init__.py +++ b/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 diff --git a/tensorflow_gnn/docs/guide/developer.md b/tensorflow_gnn/docs/guide/developer.md index 2776f7cd..4682f672 100644 --- a/tensorflow_gnn/docs/guide/developer.md +++ b/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,12 +97,16 @@ 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 @@ -107,6 +114,7 @@ example) 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*` diff --git a/tensorflow_gnn/docs/guide/keras_version.md b/tensorflow_gnn/docs/guide/keras_version.md new file mode 100644 index 00000000..d2cbc6d4 --- /dev/null +++ b/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. + + + +## 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. \ No newline at end of file diff --git a/tensorflow_gnn/docs/guide/overview.md b/tensorflow_gnn/docs/guide/overview.md index 29461d25..b49a34c3 100644 --- a/tensorflow_gnn/docs/guide/overview.md +++ b/tensorflow_gnn/docs/guide/overview.md @@ -66,6 +66,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