compatiable with both set and not set TF_USE_LEGACY_KERAS

.github/workflows/make_wheel_macOS_arm64.sh

@@ -36,7 +36,9 @@ if [[ "$TF_VERSION" == "2.11.0" && ( "$PYTHON_VERSION" == "3.9" || "$PYTHON_VERS
 fi
 
 if [[ "$TF_VERSION" =~ ^2\.(16|17)\.[0-9]+$ ]] ; then
-  export TF_USE_LEGACY_KERAS=1
+  # 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

.github/workflows/make_wheel_macOS_x86.sh

@@ -23,6 +23,7 @@ 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

.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,20 +43,14 @@ 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', '2.16.2']
+        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.7'
           - tf-version: '2.16.2'
             py-version: '3.8'
           - os: 'macos-14'
@@ -65,8 +59,6 @@ jobs:
             cpu: 'arm64'
           - os: 'macos-12'
             cpu: 'arm64'
-          - py-version: '3.7'
-            cpu: 'arm64'
           - py-version: '3.8'
             cpu: 'arm64'
       fail-fast: false
@@ -91,11 +83,6 @@ jobs:
       - uses: actions/setup-python@v2
         with:
           python-version: ${{ matrix.py-version }}
-      - if: matrix.tf-version > '2.15.1'
-        shell: bash
-        run: |
-          export TF_USE_LEGACY_KERAS=1
-          pip install --default-timeout=1000 tensorflow-estimator tf-keras==2.16.0
       - name: Setup Bazel
         # Ubuntu bazel is run inside of the docker image
         if: matrix.os != 'ubuntu-20.04'
@@ -110,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
@@ -131,19 +118,13 @@ 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', '2.16.2']
+        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.7'
           - tf-version: '2.16.2'
             py-version: '3.8'
       fail-fast: false
@@ -168,11 +149,6 @@ jobs:
       - uses: actions/setup-python@v2
         with:
           python-version: ${{ matrix.py-version }}
-      - if: matrix.tf-version > '2.15.1'
-        shell: bash
-        run: |
-          export TF_USE_LEGACY_KERAS=1
-          pip install --default-timeout=1000 tf-keras==2.16.0 tensorflow-estimator
       - name: Setup Bazel
         # Ubuntu bazel is run inside of the docker image
         if: matrix.os != 'ubuntu-20.04'
@@ -187,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
@@ -209,30 +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']
+        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.7'
           - 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

configure.py

@@ -135,7 +135,7 @@ def get_tf_version_integer():
     2.4.1 get 2041
     2.6.3 get 2063
     2.8.3 get 2083
-     get 2151
+    2.15.1 get 2151
 
   The 4-digits-string will be passed to C macro to discriminate different
   Tensorflow versions. 

tensorflow_recommenders_addons/dynamic_embedding/python/keras/layers/dynamic_layer_normalization.py

@@ -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):

tensorflow_recommenders_addons/dynamic_embedding/python/keras/layers/dynamic_layer_normalization_test.py

@@ -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]

tensorflow_recommenders_addons/dynamic_embedding/python/keras/layers/embedding.py

@@ -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,