tensorflow · rhdong · Sep 17, 2024 · Sep 9, 2024 · Sep 10, 2024
diff --git a/.github/workflows/make_wheel_Linux_x86.sh b/.github/workflows/make_wheel_Linux_x86.sh
@@ -12,6 +12,7 @@ if [[ "$TF_VERSION" =~ ^2\.(16)\.[0-9]+$ ]] ; then
   export BUILD_IMAGE="tfra/nosla-cuda12.3-cudnn8.9-ubuntu20.04-manylinux2014-python$PY_VERSION"
   export TF_CUDA_VERSION="12.3"
   export TF_CUDNN_VERSION="8.9"
+  export TF_USE_LEGACY_KERAS=1
 elif [[ "$TF_VERSION" =~ ^2\.(15)\.[0-9]+$ ]] ; then
   export BUILD_IMAGE="tfra/nosla-cuda12.2-cudnn8.9-ubuntu20.04-manylinux2014-python$PY_VERSION"
   export TF_CUDA_VERSION="12.2"

diff --git a/.github/workflows/make_wheel_Windows_x86.sh b/.github/workflows/make_wheel_Windows_x86.sh
@@ -17,6 +17,10 @@ fi
 
 python -m pip install --upgrade protobuf~=$PROTOBUF_VERSION
 
+if [[ "$TF_VERSION" =~ ^2\.(16|17)\.[0-9]+$ ]] ; then
+  export TF_USE_LEGACY_KERAS=1
+fi
+
 bash ./tools/testing/build_and_run_tests.sh
 
 python configure.py

diff --git a/.github/workflows/make_wheel_macOS_arm64.sh b/.github/workflows/make_wheel_macOS_arm64.sh
@@ -35,6 +35,12 @@ if [[ "$TF_VERSION" == "2.11.0" && ( "$PYTHON_VERSION" == "3.9" || "$PYTHON_VERS
   python -m pip install numpy==1.26.4 --force-reinstall
 fi
 
+if [[ "$TF_VERSION" =~ ^2\.(16|17)\.[0-9]+$ ]] ; then
+  # comment out the following line to test compatibility with TensorFlow 2.16 without setting TF_USE_LEGACY_KERAS=1
+  # export TF_USE_LEGACY_KERAS=1
+  python -m pip install tf-keras==2.16.0
+fi
+
 python configure.py
 
 bazel build \

diff --git a/.github/workflows/make_wheel_macOS_x86.sh b/.github/workflows/make_wheel_macOS_x86.sh
@@ -21,6 +21,11 @@ if [[ "$TF_VERSION" == "2.11.0" && ( "$PYTHON_VERSION" == "3.9" || "$PYTHON_VERS
   python -m pip install numpy==1.26.4 --force-reinstall
 fi
 
+if [[ "$TF_VERSION" =~ ^2\.(16|17)\.[0-9]+$ ]] ; then
+  export TF_USE_LEGACY_KERAS=1
+  python -m pip install tf-keras==2.16.0; \
+fi
+
 python -m pip install tensorflow-io
 python -m pip install --upgrade protobuf~=$PROTOBUF_VERSION
 

diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
@@ -15,7 +15,7 @@ on:
       - r*
 
 env:
-  MIN_PY_VERSION: '3.7'
+  MIN_PY_VERSION: '3.8'
   MAX_PY_VERSION: '3.11'
   HOROVOD_VERSION: '0.28.1'
 
@@ -43,26 +43,22 @@ jobs:
       matrix:
         # TODO: add back 'windows-latest' when it can be compiled.
         os: ['macos-14', 'macos-12', 'ubuntu-20.04']
-        py-version: ['3.7', '3.8', '3.9', '3.10', '3.11']
-        tf-version: ['2.11.0', '2.15.1']
+        py-version: ['3.8', '3.9', '3.10', '3.11']
+        tf-version: ['2.15.1', '2.16.2']
         tf-need-cuda: ['0']
         cpu: ['x86', 'arm64']
         exclude:
           # excludes cuda on macOS
-          - tf-version: '2.11.0'
-            py-version: '3.11'
-          - tf-version: '2.15.1'
-            py-version: '3.7'
           - tf-version: '2.15.1'
             py-version: '3.8'
+          - tf-version: '2.16.2'
+            py-version: '3.8'
           - os: 'macos-14'
             cpu: 'x86'
           - os: 'ubuntu-20.04'
             cpu: 'arm64'
           - os: 'macos-12'
             cpu: 'arm64'
-          - py-version: '3.7'
-            cpu: 'arm64'
           - py-version: '3.8'
             cpu: 'arm64'
       fail-fast: false
@@ -77,7 +73,7 @@ jobs:
           script: |
             const commit_details = await github.git.getCommit({owner: context.repo.owner, repo: context.repo.repo, commit_sha: context.sha});
             return commit_details.data.author.date
-      - if: matrix.tf-version != '2.15.1'
+      - if: matrix.tf-version < '2.15.1'
         shell: bash
         run: echo "SKIP_CUSTOM_OP_TESTS=--skip-custom-ops" >> $GITHUB_ENV
       - if: github.event_name == 'push'
@@ -101,7 +97,7 @@ jobs:
           CPU: ${{ matrix.cpu }}
         shell: bash
         run: |
-          if [[ "$TF_VERSION" =~ ^2\.(11|12|13|14|15)\.[0-9]$ ]] ; then
+          if [[ "$TF_VERSION" =~ ^2\.(11|12|13|14|15|16)\.[0-9]$ ]] ; then
             export HOROVOD_VERSION="0.28.1"
           fi
           bash .github/workflows/make_wheel_${OS}_${CPU}.sh
@@ -111,7 +107,7 @@ jobs:
           cond: ${{ matrix.tf-need-cuda == '1' }}
           if_true: "gpu"
           if_false: "cpu"
-      - uses: actions/upload-artifact@v1
+      - uses: actions/upload-artifact@v4
         with:
           name: ${{ runner.os }}-${{ matrix.py-version }}-tf${{ matrix.tf-version }}-${{ steps.device.outputs.value }}-${{ matrix.cpu }}-wheel
           path: wheelhouse
@@ -122,17 +118,15 @@ jobs:
       matrix:
         # TODO: add back 'windows-latest' when it can be compiled.
         os: ['ubuntu-20.04']
-        py-version: ['3.7', '3.8', '3.9', '3.10', '3.11']
-        tf-version: ['2.11.0', '2.15.1']
+        py-version: ['3.8', '3.9', '3.10', '3.11']
+        tf-version: ['2.15.1', '2.16.2']
         tf-need-cuda: ['1']
         cpu: ['x86']
         exclude:
-          - tf-version: '2.11.0'
-            py-version: '3.11'
-          - tf-version: '2.15.1'
-            py-version: '3.7'
           - tf-version: '2.15.1'
             py-version: '3.8'
+          - tf-version: '2.16.2'
+            py-version: '3.8'
       fail-fast: false
     runs-on: [self-hosted, Linux, X64]
     steps:
@@ -145,7 +139,7 @@ jobs:
           script: |
             const commit_details = await github.git.getCommit({owner: context.repo.owner, repo: context.repo.repo, commit_sha: context.sha});
             return commit_details.data.author.date
-      - if: matrix.tf-version != '2.15.1'
+      - if: matrix.tf-version < '2.15.1'
         shell: bash
         run: echo "SKIP_CUSTOM_OP_TESTS=--skip-custom-ops" >> $GITHUB_ENV
       - if: github.event_name == 'push'
@@ -169,7 +163,7 @@ jobs:
           CPU: ${{ matrix.cpu }}
         shell: bash
         run: |
-          if [[ "$TF_VERSION" =~ ^2\.(11|12|13|14|15)\.[0-9]$ ]] ; then
+          if [[ "$TF_VERSION" =~ ^2\.(11|12|13|14|15|16)\.[0-9]$ ]] ; then
             export HOROVOD_VERSION="0.28.1"
           fi
           bash .github/workflows/make_wheel_${OS}_${CPU}.sh
@@ -179,7 +173,7 @@ jobs:
           cond: ${{ matrix.tf-need-cuda == '1' }}
           if_true: "gpu"
           if_false: "cpu"
-      - uses: actions/upload-artifact@v1
+      - uses: actions/upload-artifact@v4
         with:
           name: ${{ runner.os }}-${{ matrix.py-version }}-tf${{ matrix.tf-version }}-${{ steps.device.outputs.value }}-${{ matrix.cpu }}-wheel
           path: wheelhouse
@@ -191,26 +185,22 @@ jobs:
       matrix:
         # TODO: add back 'Windows' when it can be compiled.
         os: ['macOS', 'Linux']
-        py-version: ['3.7', '3.8', '3.9', '3.10', '3.11']
-        tf-version: ['2.15.1']
+        py-version: ['3.8', '3.9', '3.10', '3.11']
+        tf-version: ['2.15.1', '2.16.2']
         tf-need-cuda: ['0', '1']
         cpu: ['x86', 'arm64']
         exclude:
           # excludes cuda on macOS
           - os: 'macOS'
             tf-need-cuda: '1'
-          - tf-version: '2.11.0'
-            py-version: '3.11'
-          - tf-version: '2.15.1'
-            py-version: '3.7'
           - tf-version: '2.15.1'
             py-version: '3.8'
+          - tf-version: '2.16.2'
+            py-version: '3.8'
           - os: 'Linux'
             cpu: 'arm64'
           - os: 'Linux'
             tf-need-cuda: '0'
-          - py-version: '3.7'
-            cpu: 'arm64'
           - py-version: '3.8'
             cpu: 'arm64'
       fail-fast: false

diff --git a/configure.py b/configure.py
@@ -150,7 +150,7 @@ def get_tf_version_integer():
         '\nPlease install a TensorFlow on your compiling machine, '
         'The compiler needs to know the version of Tensorflow '
         'and get TF c++ headers according to the installed TensorFlow. '
-        '\nNote: Only TensorFlow 2.15.1 2.8.3, 2.6.3, 2.4.1, 1.15.2 are supported.'
+        '\nNote: Only TensorFlow 2.16.2 2.15.1 2.8.3, 2.6.3, 2.4.1, 1.15.2 are supported.'
     )
   try:
     major, minor, patch = version.split('.')

@@ -1,5 +1,8 @@
 import tensorflow as tf
-from tensorflow.keras.layers import LayerNormalization as TFLayerNormalization
+try:  # tf version <= 2.15
+  from tensorflow.keras.layers import LayerNormalization as TFLayerNormalization
+except:
+  from tf_keras.layers import LayerNormalization as TFLayerNormalization
 
 
 class LayerNormalization(TFLayerNormalization):

@@ -3,15 +3,20 @@
 import tensorflow as tf
 from tensorflow_recommenders_addons import dynamic_embedding as de
 
+try:
+  from tf_keras import Input, models, layers
+except:
+  from tensorflow.keras import Input, models, layers
+
 
 class DynamicLayerNormalizationTest(tf.test.TestCase):
 
   def test_dynamic_shape_support(self):
-    input_data = tf.keras.Input(shape=(None, 10), dtype=tf.float32)
+    input_data = Input(shape=(None, 10), dtype=tf.float32)
     layer = de.keras.layers.LayerNormalization()
     output = layer(input_data)
 
-    model = tf.keras.models.Model(inputs=input_data, outputs=output)
+    model = models.Model(inputs=input_data, outputs=output)
 
     np.random.seed(0)
     test_data = np.random.randn(2, 5, 10).astype(np.float32)
@@ -38,10 +43,10 @@ def test_training_with_layer_normalization(self):
               np.random.randn(num_samples) * 0.5).astype(np.float32).reshape(
                   -1, 1)
 
-    input_data = tf.keras.Input(shape=(input_dim,), dtype=tf.float32)
+    input_data = Input(shape=(input_dim,), dtype=tf.float32)
     normalized = de.keras.layers.LayerNormalization()(input_data)
-    output = tf.keras.layers.Dense(output_dim)(normalized)
-    model = tf.keras.models.Model(inputs=input_data, outputs=output)
+    output = layers.Dense(output_dim)(normalized)
+    model = models.Model(inputs=input_data, outputs=output)
 
     model.compile(optimizer='adam', loss='mean_squared_error')
     initial_weights = [layer.get_weights() for layer in model.layers]

@@ -48,6 +48,15 @@
   TrainableWrapperDistributedPolicy
 from tensorflow_recommenders_addons.dynamic_embedding.python.ops.tf_save_restore_patch import de_fs_saveable_class_names
 
+try:  # tf version >= 2.16
+  from tf_keras.layers import Layer
+  from tf_keras.initializers import RandomNormal, Zeros, serialize
+  from tf_keras import constraints
+except:
+  from tensorflow.keras.layers import Layer
+  from tensorflow.keras.initializers import RandomNormal, Zeros, serialize
+  from tensorflow.keras import constraints
+
 
 def _choose_reduce_method(combiner, sparse=False, segmented=False):
   select = 'sparse' if sparse else 'math'
@@ -93,7 +102,7 @@ def reduce_pooling(x, combiner='sum'):
     return x
 
 
-class Embedding(tf.keras.layers.Layer):
+class Embedding(Layer):
   """
   A keras style Embedding layer. The `Embedding` layer acts same like
   [tf.keras.layers.Embedding](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Embedding),
@@ -180,7 +189,7 @@ def __init__(self,
     self.embedding_size = embedding_size
     self.combiner = combiner
     if initializer is None:
-      initializer = tf.keras.initializers.RandomNormal()
+      initializer = RandomNormal()
     partitioner = kwargs.get('partitioner', devar.default_partition_fn)
     trainable = kwargs.get('trainable', True)
     self.max_norm = kwargs.get('max_norm', None)
@@ -281,10 +290,10 @@ def call(self, ids):
   def get_config(self):
     _initializer = self.params.initializer
     if _initializer is None:
-      _initializer = tf.keras.initializers.Zeros()
+      _initializer = Zeros()
     _max_norm = None
-    if isinstance(self.max_norm, tf.keras.constraints.Constraint):
-      _max_norm = tf.keras.constraints.serialize(self.max_norm)
+    if isinstance(self.max_norm, constraints.Constraint):
+      _max_norm = constraints.serialize(self.max_norm)
 
     if self.params.restrict_policy:
       _restrict_policy = self.params.restrict_policy.__class__
@@ -301,7 +310,7 @@ def get_config(self):
         'combiner':
             self.combiner,
         'initializer':
-            tf.keras.initializers.serialize(_initializer),
+            serialize(_initializer),
         'devices':
             self.params.devices if self.keep_distribution else None,
         'name':
@@ -500,10 +509,10 @@ def _pooling_by_slots(self, lookup_result, ids):
   def get_config(self):
     _initializer = self.params.initializer
     if _initializer is None:
-      _initializer = tf.keras.initializers.Zeros()
+      _initializer = Zeros()
     _max_norm = None
-    if isinstance(self.max_norm, tf.keras.constraints.Constraint):
-      _max_norm = tf.keras.constraints.serialize(self.max_norm)
+    if isinstance(self.max_norm, constraints.Constraint):
+      _max_norm = constraints.serialize(self.max_norm)
 
     config = {
         'embedding_size': self.embedding_size,
@@ -512,7 +521,7 @@ def get_config(self):
         'combiner': self.combiner,
         'key_dtype': self.params.key_dtype,
         'value_dtype': self.params.value_dtype,
-        'initializer': tf.keras.initializers.serialize(_initializer),
+        'initializer': serialize(_initializer),
         'devices': self.params.devices,
         'name': self.name,
         'trainable': self.trainable,