diff --git a/.github/workflows/make_wheel_macOS_arm64.sh b/.github/workflows/make_wheel_macOS_arm64.sh
index 36a2ea9be..3607d0e59 100644
--- a/.github/workflows/make_wheel_macOS_arm64.sh
+++ b/.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
diff --git a/.github/workflows/make_wheel_macOS_x86.sh b/.github/workflows/make_wheel_macOS_x86.sh
index ba43551d1..86eaf8d68 100644
--- a/.github/workflows/make_wheel_macOS_x86.sh
+++ b/.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
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
index d8d24e5f0..f03f50721 100644
--- 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,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,7 +185,7 @@ 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']
@@ -217,22 +193,14 @@ jobs:
           # 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
diff --git a/configure.py b/configure.py
index 4dfad3aac..47bf3a99b 100644
--- a/configure.py
+++ b/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. 
diff --git a/tensorflow_recommenders_addons/dynamic_embedding/python/keras/layers/dynamic_layer_normalization.py b/tensorflow_recommenders_addons/dynamic_embedding/python/keras/layers/dynamic_layer_normalization.py
index 26f64d6d6..7a6bb5b17 100644
--- a/tensorflow_recommenders_addons/dynamic_embedding/python/keras/layers/dynamic_layer_normalization.py
+++ b/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):
diff --git a/tensorflow_recommenders_addons/dynamic_embedding/python/keras/layers/dynamic_layer_normalization_test.py b/tensorflow_recommenders_addons/dynamic_embedding/python/keras/layers/dynamic_layer_normalization_test.py
index 59dfdfc87..afe6f8f49 100644
--- a/tensorflow_recommenders_addons/dynamic_embedding/python/keras/layers/dynamic_layer_normalization_test.py
+++ b/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]
diff --git a/tensorflow_recommenders_addons/dynamic_embedding/python/keras/layers/embedding.py b/tensorflow_recommenders_addons/dynamic_embedding/python/keras/layers/embedding.py
index 8e82bdef3..736f16ca9 100644
--- a/tensorflow_recommenders_addons/dynamic_embedding/python/keras/layers/embedding.py
+++ b/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,
diff --git a/tensorflow_recommenders_addons/dynamic_embedding/python/keras/layers/embedding_test.py b/tensorflow_recommenders_addons/dynamic_embedding/python/keras/layers/embedding_test.py
index c5087a8e6..f1d97b80b 100644
--- a/tensorflow_recommenders_addons/dynamic_embedding/python/keras/layers/embedding_test.py
+++ b/tensorflow_recommenders_addons/dynamic_embedding/python/keras/layers/embedding_test.py
@@ -39,22 +39,32 @@
 from tensorflow.python.platform import test
 from tensorflow.python.training import adam
 try:
-  from tensorflow.keras.optimizers.legacy import Adam
+  from tf_keras.optimizers import Adam
 except:
-  from tensorflow.keras.optimizers import Adam
+  try:
+    from tensorflow.keras.optimizers import Adam
+  except:
+    from tensorflow.keras.optimizers.legacy import Adam
+
+try:
+  from tf_keras import layers, Sequential, optimizers, Model, backend
+  from tf_keras.initializers import RandomNormal, RandomUniform, Ones, Zeros
+except:
+  from tensorflow.keras import layers, Sequential, optimizers, Model, backend
+  from tensorflow.keras.initializers import RandomNormal, RandomUniform, Ones, Zeros
 
 tf.config.set_soft_device_placement(True)
 
 
 def get_sequential_model(emb_t, *args, **kwargs):
-  l0 = tf.keras.layers.InputLayer(input_shape=(None,), dtype=dtypes.int64)
+  l0 = layers.InputLayer(input_shape=(None,), dtype=dtypes.int64)
   l1 = emb_t(*args, **kwargs)
-  l2 = tf.keras.layers.Dense(8, 'relu')
-  l3 = tf.keras.layers.Dense(1, 'sigmoid')
+  l2 = layers.Dense(8, 'relu')
+  l3 = layers.Dense(1, 'sigmoid')
   if emb_t == de.keras.layers.Embedding:
-    model = tf.keras.Sequential([l0, l1, l2, l3])
+    model = Sequential([l0, l1, l2, l3])
   elif emb_t == de.keras.layers.FieldWiseEmbedding:
-    model = tf.keras.Sequential([l0, l1, tf.keras.layers.Flatten(), l2, l3])
+    model = Sequential([l0, l1, layers.Flatten(), l2, l3])
   else:
     raise TypeError('Unsupported embedding layer {}'.format(emb_t))
   return model
@@ -80,10 +90,7 @@ def test_create(self):
     value_dtypes = [dtypes.float32, dtypes.float64]
     if test_util.is_gpu_available():
       value_dtypes = [dtypes.float32]
-    initializers = [
-        tf.keras.initializers.RandomNormal(),
-        tf.keras.initializers.RandomUniform()
-    ]
+    initializers = [RandomNormal(), RandomUniform()]
     trainable_options = [True, False]
     bp_options = [True, False]
     restrict_policies = [
@@ -115,13 +122,13 @@ def test_create(self):
   def test_forward(self):
     if not context.executing_eagerly():
       self.skipTest('Only test in eager mode')
-    de_init = tf.keras.initializers.RandomNormal(seed=0)
-    dense_init = tf.keras.initializers.Ones()
+    de_init = RandomNormal(seed=0)
+    dense_init = Ones()
     de_layer = de.keras.layers.Embedding(4, initializer=de_init, name='ve820')
-    tf_layer = tf.keras.layers.Embedding(1000,
-                                         4,
-                                         embeddings_initializer=dense_init,
-                                         name='mt047')
+    tf_layer = layers.Embedding(1000,
+                                4,
+                                embeddings_initializer=dense_init,
+                                name='mt047')
 
     with self.session(use_gpu=test_util.is_gpu_available(),
                       config=default_config):
@@ -137,7 +144,7 @@ def test_forward(self):
   def test_backward(self):
     if not context.executing_eagerly():
       self.skipTest('Only test in eager mode')
-    init = tf.keras.initializers.RandomNormal(seed=0)
+    init = RandomNormal(seed=0)
     model = get_sequential_model(de.keras.layers.Embedding,
                                  4,
                                  initializer=init,
@@ -160,13 +167,13 @@ def test_backward(self):
   def test_backward_adagrad(self):
     if not context.executing_eagerly():
       self.skipTest('Only test in eager mode')
-    init = tf.keras.initializers.RandomNormal(seed=0)
+    init = RandomNormal(seed=0)
     model = get_sequential_model(de.keras.layers.Embedding,
                                  4,
                                  initializer=init,
                                  bp_v2=False,
                                  name='go582')
-    optmz = tf.keras.optimizers.Adagrad(1E-4)
+    optmz = optimizers.Adagrad(1E-4)
     optmz = de.DynamicEmbeddingOptimizer(optmz)
     emb_layer = model.layers[0]
     model.compile(optimizer=optmz, loss='binary_crossentropy')
@@ -183,7 +190,7 @@ def test_backward_adagrad(self):
   def test_backward_bp_v2(self):
     if not context.executing_eagerly():
       self.skipTest('Only test in eager mode')
-    init = tf.keras.initializers.RandomNormal(seed=0)
+    init = RandomNormal(seed=0)
     model = get_sequential_model(de.keras.layers.Embedding,
                                  4,
                                  initializer=init,
@@ -210,16 +217,16 @@ def test_keras_save_load_weights(self):
     save_path = os.path.join(tempfile.mkdtemp(prefix=save_dir), "hash")
 
     def model_fn(table_device):
-      input_tensor = tf.keras.layers.Input(shape=(1,), dtype=tf.int64)
+      input_tensor = layers.Input(shape=(1,), dtype=tf.int64)
       embedding_out = de.keras.layers.Embedding(
           embedding_size=1,
           key_dtype=tf.int64,
           value_dtype=tf.float32,
-          initializer=tf.keras.initializers.RandomNormal(),
+          initializer=RandomNormal(),
           devices=table_device,
           name='test_keras_save_restore',
       )(input_tensor)
-      model = tf.keras.Model(inputs=input_tensor, outputs=embedding_out)
+      model = Model(inputs=input_tensor, outputs=embedding_out)
       optimizer = Adam(learning_rate=1E-4, amsgrad=False)
       optimizer = de.DynamicEmbeddingOptimizer(optimizer)
       model.compile(optimizer=optimizer)
@@ -236,7 +243,7 @@ def model_fn(table_device):
     )
     options = tf.saved_model.SaveOptions(namespace_whitelist=['TFRA'])
     model.save(save_path, options=options)
-    tf.keras.backend.clear_session()
+    backend.clear_session()
     del model
     model = model_fn(table_device_)
     model.load_weights(save_path).expect_partial()
@@ -264,14 +271,14 @@ def test_keras_save_load_weights_file_system(self):
     save_path = os.path.join(tempfile.mkdtemp(prefix=save_dir), "hash")
 
     def model_fn(table_devices):
-      input_tensor = tf.keras.layers.Input(shape=(1,), dtype=tf.int64)
+      input_tensor = layers.Input(shape=(1,), dtype=tf.int64)
       embedding_outs = []
       for t in range(2):
         embedding_out = de.keras.layers.Embedding(
             embedding_size=1,
             key_dtype=tf.int64,
             value_dtype=tf.float32,
-            initializer=tf.keras.initializers.RandomNormal(),
+            initializer=RandomNormal(),
             devices=table_devices,
             name=f'test_keras_save_restore_{t}',
             kv_creator=de.CuckooHashTableCreator(
@@ -281,12 +288,12 @@ def model_fn(table_devices):
           embedding_size=1,
           key_dtype=tf.int64,
           value_dtype=tf.float32,
-          initializer=tf.keras.initializers.RandomNormal(),
+          initializer=RandomNormal(),
           devices=table_devices,
           name='test_keras_save_restore_normal')(input_tensor)
       embedding_outs.append(normal_embedding_out)
       concat = tf.concat(embedding_outs, axis=0)
-      model = tf.keras.Model(inputs=input_tensor, outputs=concat)
+      model = Model(inputs=input_tensor, outputs=concat)
       optimizer = Adam(learning_rate=1E-4, amsgrad=False)
       optimizer = de.DynamicEmbeddingOptimizer(optimizer)
       model.compile(optimizer=optimizer)
@@ -309,7 +316,7 @@ def model_fn(table_devices):
       )
     options = tf.saved_model.SaveOptions(namespace_whitelist=['TFRA'])
     model.save(save_path, options=options)
-    tf.keras.backend.clear_session()
+    backend.clear_session()
     del model
     model = model_fn(table_devices_)
     model.load_weights(save_path).expect_partial()
@@ -322,7 +329,7 @@ def model_fn(table_devices):
       self.assertAllEqual(test_values, np.sort(values, axis=0))
 
     # test expand shards number
-    tf.keras.backend.clear_session()
+    backend.clear_session()
     del model
     shard_num = 5
     table_devices_ = table_device * shard_num
@@ -337,7 +344,7 @@ def model_fn(table_devices):
       self.assertAllEqual(test_values, np.sort(values, axis=0))
 
     # test contracte shards number
-    tf.keras.backend.clear_session()
+    backend.clear_session()
     del model
     shard_num = 2
     table_devices_ = table_device * shard_num
@@ -352,7 +359,7 @@ def model_fn(table_devices):
       self.assertAllEqual(test_values, np.sort(values, axis=0))
 
     # test load all into one shard
-    tf.keras.backend.clear_session()
+    backend.clear_session()
     del model
     shard_num = 1
     table_devices_ = table_device * shard_num
@@ -377,14 +384,14 @@ def model_fn(proc_size, proc_rank):
       table_device = ['/device:CPU:0']
       if test_util.is_gpu_available():
         table_device = ['/device:GPU:0']
-      input_tensor = tf.keras.layers.Input(shape=(1,), dtype=tf.int64)
+      input_tensor = layers.Input(shape=(1,), dtype=tf.int64)
       embedding_outs = []
       for t in range(2):
         embedding_out = de.keras.layers.Embedding(
             embedding_size=1,
             key_dtype=tf.int64,
             value_dtype=tf.float32,
-            initializer=tf.keras.initializers.RandomNormal(),
+            initializer=RandomNormal(),
             devices=table_device,
             name=f'test_keras_save_restore_{t}',
             kv_creator=de.CuckooHashTableCreator(saver=de.FileSystemSaver(
@@ -394,12 +401,12 @@ def model_fn(proc_size, proc_rank):
           embedding_size=1,
           key_dtype=tf.int64,
           value_dtype=tf.float32,
-          initializer=tf.keras.initializers.RandomNormal(),
+          initializer=RandomNormal(),
           devices=table_device,
           name='test_keras_save_restore_normal')(input_tensor)
       embedding_outs.append(normal_embedding_out)
       concat = tf.concat(embedding_outs, axis=0)
-      model = tf.keras.Model(inputs=input_tensor, outputs=concat)
+      model = Model(inputs=input_tensor, outputs=concat)
       optimizer = Adam(learning_rate=1E-4, amsgrad=False)
       optimizer = de.DynamicEmbeddingOptimizer(optimizer)
       model.compile(optimizer=optimizer)
@@ -417,7 +424,7 @@ def model_fn(proc_size, proc_rank):
     keys_shard_size = int(test_size / proc_size)
     models = []
     for i in range(proc_size):
-      tf.keras.backend.clear_session()
+      backend.clear_session()
       models.append(model_fn(proc_size, i))
       for t in range(2):
         params_ = models[i].get_layer(f'test_keras_save_restore_{t}').params
@@ -437,11 +444,11 @@ def model_fn(proc_size, proc_rank):
           models[i].get_layer(
               f'test_keras_save_restore_{t}').params.save_to_file_system(
                   dirpath=de_dir, proc_size=proc_size, proc_rank=i)
-    tf.keras.backend.clear_session()
+    backend.clear_session()
     for i in range(proc_size):
       del models[0]
     for i in range(proc_size):
-      tf.keras.backend.clear_session()
+      backend.clear_session()
       models.append(model_fn(proc_size, i))
       models[i].load_weights(save_path).expect_partial()
     for t in range(2):
@@ -464,7 +471,7 @@ def model_fn(proc_size, proc_rank):
     keys_shard_size = int(test_size / proc_size)
     models = []
     for i in range(proc_size):
-      tf.keras.backend.clear_session()
+      backend.clear_session()
       models.append(model_fn(proc_size, i))
       for t in range(2):
         params_ = models[i].get_layer(f'test_keras_save_restore_{t}').params
@@ -484,12 +491,12 @@ def model_fn(proc_size, proc_rank):
           models[i].get_layer(
               f'test_keras_save_restore_{t}').params.save_to_file_system(
                   dirpath=de_dir, proc_size=proc_size, proc_rank=i)
-    tf.keras.backend.clear_session()
+    backend.clear_session()
     for i in range(proc_size):
       del models[0]
     proc_size = 5
     for i in range(proc_size):
-      tf.keras.backend.clear_session()
+      backend.clear_session()
       models.append(model_fn(proc_size, i))
       models[i].load_weights(save_path).expect_partial()
     for t in range(2):
@@ -512,7 +519,7 @@ def model_fn(proc_size, proc_rank):
     keys_shard_size = int(test_size / proc_size)
     models = []
     for i in range(proc_size):
-      tf.keras.backend.clear_session()
+      backend.clear_session()
       models.append(model_fn(proc_size, i))
       for t in range(2):
         params_ = models[i].get_layer(f'test_keras_save_restore_{t}').params
@@ -532,11 +539,11 @@ def model_fn(proc_size, proc_rank):
           models[i].get_layer(
               f'test_keras_save_restore_{t}').params.save_to_file_system(
                   dirpath=de_dir, proc_size=proc_size, proc_rank=i)
-    tf.keras.backend.clear_session()
+    backend.clear_session()
     for i in range(proc_size):
       del models[0]
     for i in range(proc_size):
-      tf.keras.backend.clear_session()
+      backend.clear_session()
       models.append(model_fn(proc_size, i))
       models[i].load_weights(save_path).expect_partial()
     for t in range(2):
@@ -559,7 +566,7 @@ def model_fn(proc_size, proc_rank):
     keys_shard_size = int(test_size / proc_size)
     models = []
     for i in range(proc_size):
-      tf.keras.backend.clear_session()
+      backend.clear_session()
       models.append(model_fn(proc_size, i))
       for t in range(2):
         params_ = models[i].get_layer(f'test_keras_save_restore_{t}').params
@@ -579,12 +586,12 @@ def model_fn(proc_size, proc_rank):
           models[i].get_layer(
               f'test_keras_save_restore_{t}').params.save_to_file_system(
                   dirpath=de_dir, proc_size=proc_size, proc_rank=i)
-    tf.keras.backend.clear_session()
+    backend.clear_session()
     for i in range(proc_size):
       del models[0]
     proc_size = 2
     for i in range(proc_size):
-      tf.keras.backend.clear_session()
+      backend.clear_session()
       models.append(model_fn(proc_size, i))
       models[i].load_weights(save_path).expect_partial()
     for t in range(2):
@@ -607,7 +614,7 @@ def model_fn(proc_size, proc_rank):
     keys_shard_size = int(test_size / proc_size)
     models = []
     for i in range(proc_size):
-      tf.keras.backend.clear_session()
+      backend.clear_session()
       models.append(model_fn(proc_size, i))
       for t in range(2):
         params_ = models[i].get_layer(f'test_keras_save_restore_{t}').params
@@ -627,11 +634,11 @@ def model_fn(proc_size, proc_rank):
           models[i].get_layer(
               f'test_keras_save_restore_{t}').params.save_to_file_system(
                   dirpath=de_dir, proc_size=proc_size, proc_rank=i)
-    tf.keras.backend.clear_session()
+    backend.clear_session()
     for i in range(proc_size):
       del models[0]
     for i in range(proc_size):
-      tf.keras.backend.clear_session()
+      backend.clear_session()
       models.append(model_fn(proc_size, i))
       models[i].load_weights(save_path).expect_partial()
     for t in range(2):
@@ -654,7 +661,7 @@ def model_fn(proc_size, proc_rank):
     keys_shard_size = int(test_size / proc_size)
     models = []
     for i in range(proc_size):
-      tf.keras.backend.clear_session()
+      backend.clear_session()
       models.append(model_fn(proc_size, i))
       for t in range(2):
         params_ = models[i].get_layer(f'test_keras_save_restore_{t}').params
@@ -674,7 +681,7 @@ def model_fn(proc_size, proc_rank):
           models[i].get_layer(
               f'test_keras_save_restore_{t}').params.save_to_file_system(
                   dirpath=de_dir, proc_size=proc_size, proc_rank=i)
-    tf.keras.backend.clear_session()
+    backend.clear_session()
     for i in range(proc_size):
       del models[0]
     proc_size = 1
@@ -682,7 +689,7 @@ def model_fn(proc_size, proc_rank):
     total_keys = []
     total_values = []
     for i in range(proc_size):
-      tf.keras.backend.clear_session()
+      backend.clear_session()
       models.append(model_fn(proc_size, i))
       models[i].load_weights(save_path).expect_partial()
     for t in range(2):
@@ -707,17 +714,17 @@ class SquashedEmbeddingLayerTest(test.TestCase):
   def test_forward(self):
     if not context.executing_eagerly():
       self.skipTest('Only test in eager mode')
-    init = tf.keras.initializers.Zeros()
+    init = Zeros()
     de_layer = de.keras.layers.SquashedEmbedding(2,
                                                  initializer=init,
                                                  key_dtype=dtypes.int64,
                                                  value_dtype=dtypes.float32,
                                                  name='tr423')
-    dense_init = tf.keras.initializers.Ones()
-    tf_layer = tf.keras.layers.Embedding(100,
-                                         2,
-                                         embeddings_initializer=dense_init,
-                                         name='mt047')
+    dense_init = Ones()
+    tf_layer = layers.Embedding(100,
+                                2,
+                                embeddings_initializer=dense_init,
+                                name='mt047')
 
     preset_ids = constant_op.constant([3, 0, 1], dtype=dtypes.int64)
     preset_values = constant_op.constant([[1, 1], [1, 1], [1, 1]],
@@ -746,10 +753,7 @@ def test_create(self):
     value_dtypes = [dtypes.float32, dtypes.float64]
     if test_util.is_gpu_available():
       value_dtypes = [dtypes.float32]
-    initializers = [
-        tf.keras.initializers.RandomNormal(),
-        tf.keras.initializers.RandomUniform()
-    ]
+    initializers = [RandomNormal(), RandomUniform()]
     trainable_options = [True, False]
     bp_options = [True, False]
     restrict_policies = [
@@ -791,7 +795,7 @@ def test_forward(self):
     if not context.executing_eagerly():
       self.skipTest('Only test in eager mode')
 
-    init = tf.keras.initializers.RandomNormal(seed=0)
+    init = RandomNormal(seed=0)
     ids = math_ops.range(200, dtype=dtypes.int64)
     ids = tf.reshape(ids, (25, 8))
 
@@ -803,10 +807,10 @@ def slot_map_fn(x):
                                                   slot_map_fn,
                                                   initializer=init,
                                                   name='fr010')
-    tf_layer = tf.keras.layers.Embedding(1000,
-                                         4,
-                                         embeddings_initializer=init,
-                                         name='xz774')
+    tf_layer = layers.Embedding(1000,
+                                4,
+                                embeddings_initializer=init,
+                                name='xz774')
     with self.session(use_gpu=test_util.is_gpu_available(),
                       config=default_config):
       self.evaluate(variables.global_variables_initializer())
@@ -821,7 +825,7 @@ def test_backward(self):
     def slot_map_fn(x):
       return tf.math.floormod(x, 2)
 
-    init = tf.keras.initializers.RandomNormal(seed=0)
+    init = RandomNormal(seed=0)
     model = get_sequential_model(de.keras.layers.FieldWiseEmbedding,
                                  4,
                                  2,
@@ -851,7 +855,7 @@ def test_sequential_model_save_and_load_weights(self):
     def slot_map_fn(x):
       return tf.math.floormod(x, 2)
 
-    init = tf.keras.initializers.RandomNormal(seed=0)
+    init = RandomNormal(seed=0)
     model = get_sequential_model(de.keras.layers.FieldWiseEmbedding,
                                  4,
                                  2,
@@ -880,7 +884,7 @@ def slot_map_fn(x):
     save_dir = tempfile.mkdtemp(prefix='/tmp/')
     options = tf.saved_model.SaveOptions(namespace_whitelist=['TFRA'])
     model.save(save_dir, signatures=None, options=options)
-    copied_init = tf.keras.initializers.RandomNormal(seed=0)
+    copied_init = RandomNormal(seed=0)
     new_model = get_sequential_model(de.keras.layers.FieldWiseEmbedding,
                                      4,
                                      2,
@@ -928,23 +932,22 @@ def test_model_save_and_load(self):
     def slot_map_fn(x):
       return tf.math.floormod(x, 2)
 
-    init = tf.keras.initializers.RandomNormal(seed=0)
+    init = RandomNormal(seed=0)
 
-    class MyModel(tf.keras.Model):
+    class MyModel(Model):
 
       def __init__(self):
         super(MyModel, self).__init__()
-        self.l0 = tf.keras.layers.InputLayer(input_shape=(None,),
-                                             dtype=tf.int64)
+        self.l0 = layers.InputLayer(input_shape=(None,), dtype=tf.int64)
         self.l1 = de.keras.layers.FieldWiseEmbedding(4,
                                                      2,
                                                      slot_map_fn,
                                                      bp_v2=False,
                                                      initializer=init,
                                                      name='sl337')
-        self.l2 = tf.keras.layers.Flatten()
-        self.l3 = tf.keras.layers.Dense(32, 'relu')
-        self.l4 = tf.keras.layers.Dense(1, 'sigmoid')
+        self.l2 = layers.Flatten()
+        self.l3 = layers.Dense(32, 'relu')
+        self.l4 = layers.Dense(1, 'sigmoid')
 
       def call(self, x):
         return self.l4(self.l3(self.l2(self.l1(self.l0(x)))))
diff --git a/tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/dynamic_embedding_ops_test.py b/tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/dynamic_embedding_ops_test.py
index 8b991226a..cb60c3063 100644
--- a/tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/dynamic_embedding_ops_test.py
+++ b/tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/dynamic_embedding_ops_test.py
@@ -67,10 +67,16 @@
 from tensorflow.python.training import device_setter
 from tensorflow.python.training import server_lib
 from tensorflow.python.util import compat
+
 try:
-  from tensorflow.keras.legacy.optimizers import Adam
+  from tf_keras.initializers import Zeros
+  from tf_keras.optimizers import Adam
 except:
-  from tensorflow.keras.optimizers import Adam
+  from tensorflow.keras.initializers import Zeros
+  try:
+    from tensorflow.keras.optimizers import Adam
+  except:
+    from tensorflow.keras.legacy.optimizers import Adam
 
 
 # pylint: disable=missing-class-docstring
@@ -1415,9 +1421,8 @@ def _create_input_and_params(self,
     labels = array_ops.zeros((batch_size,), dtype=dtypes.float32)
     devar = de.get_variable(name + '/dynamic_embedding',
                             dim=embedding_size,
-                            initializer=tf.keras.initializers.Zeros())
-    tfvar = tf.Variable(tf.keras.initializers.Zeros()((nids, embedding_size),
-                                                      dtype=tf.float32))
+                            initializer=Zeros())
+    tfvar = tf.Variable(Zeros()((nids, embedding_size), dtype=tf.float32))
     return ids, labels, devar, tfvar
 
   def _loss_fn(self, params, ids, labels):
diff --git a/tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/dynamic_embedding_variable_test.py b/tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/dynamic_embedding_variable_test.py
index 79524b6e1..0c9f736ba 100644
--- a/tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/dynamic_embedding_variable_test.py
+++ b/tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/dynamic_embedding_variable_test.py
@@ -30,7 +30,6 @@
 from tensorflow_recommenders_addons.utils.check_platform import is_macos, is_arm64
 
 from tensorflow.core.protobuf import config_pb2
-from tensorflow.keras import layers
 from tensorflow.python.client import session
 from tensorflow.python.eager import context
 from tensorflow.python.framework import constant_op
@@ -66,9 +65,9 @@
 try:  # tf version <= 2.15
   from tensorflow_estimator.python.estimator import estimator
   from tensorflow_estimator.python.estimator import estimator_lib
+  from tensorflow.keras import layers
 except:
-  # do nothing
-  pass
+  from tf_keras import layers
 
 try:  # The data_structures has been moved to the new package in tf 2.11
   from tensorflow.python.trackable import data_structures
diff --git a/tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/horovod_sync_train_test.py b/tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/horovod_sync_train_test.py
index de768e12a..1f94f48e7 100644
--- a/tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/horovod_sync_train_test.py
+++ b/tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/horovod_sync_train_test.py
@@ -38,10 +38,19 @@
 from tensorflow.python.training import checkpoint_management
 from tensorflow.python.training import monitored_session
 from tensorflow.python.training import training_util
+
 try:
-  from tensorflow.keras.optimizers.legacy import Adam
+  from tf_keras import layers, Sequential, models, backend
+  from tf_keras.initializers import Zeros
+  from tf_keras.optimizers import Adam
 except:
-  from tensorflow.keras.optimizers import Adam
+  from tensorflow.keras import layers, Sequential, models, backend
+  from tensorflow.keras.initializers import Zeros
+  try:
+    from tensorflow.keras.optimizers import Adam
+  except:
+    from tensorflow.keras.legacy.optimizers import Adam
+
 from tensorflow_recommenders_addons.utils.check_platform import is_macos, is_arm64
 
 default_config = config_pb2.ConfigProto(
@@ -50,14 +59,14 @@
 
 
 def get_emb_sequential_model(emb_t, opt, *args, **kwargs):
-  l0 = tf.keras.layers.InputLayer(input_shape=(None,), dtype=dtypes.int64)
+  l0 = layers.InputLayer(input_shape=(None,), dtype=dtypes.int64)
   l1 = emb_t(*args, **kwargs)
-  l2 = tf.keras.layers.Dense(8, 'relu', kernel_initializer='zeros')
-  l3 = tf.keras.layers.Dense(1, 'sigmoid', kernel_initializer='zeros')
-  if emb_t == tf.keras.layers.Embedding:
-    model = tf.keras.Sequential([l0, l1, l2, l3])
+  l2 = layers.Dense(8, 'relu', kernel_initializer='zeros')
+  l3 = layers.Dense(1, 'sigmoid', kernel_initializer='zeros')
+  if emb_t == layers.Embedding:
+    model = Sequential([l0, l1, l2, l3])
   elif emb_t == de.keras.layers.HvdAllToAllEmbedding:
-    model = tf.keras.Sequential([l0, l1, l2, l3])
+    model = Sequential([l0, l1, l2, l3])
   else:
     raise TypeError('Unsupported embedding layer {}'.format(emb_t))
   model.compile(optimizer=opt, loss='mean_absolute_error')
@@ -80,7 +89,7 @@ def test_adam_minimize_trainable(self):
     base_opt = adam.AdamOptimizer(1.0)
     test_opt = adam.AdamOptimizer(1.0)
     self.common_minimize_trainable_v1(base_opt, test_opt, name="adam")
-    tf.keras.backend.clear_session()
+    backend.clear_session()
     keras_base_opt = Adam(1.0)
     keras_test_opt = Adam(1.0)
     self.common_minimize_trainable_v2(keras_base_opt,
@@ -281,7 +290,7 @@ def common_all_to_all_embedding_trainable_v2(self, base_opt, test_opt, name):
 
     base_opt = de.DynamicEmbeddingOptimizer(base_opt, synchronous=True)
     test_opt = hvd.DistributedOptimizer(test_opt)
-    init = tf.keras.initializers.Zeros()
+    init = Zeros()
     kv_creator = de.CuckooHashTableCreator(
         saver=de.FileSystemSaver(proc_size=hvd.size(), proc_rank=hvd.rank()))
     batch_size = 8
@@ -303,7 +312,7 @@ def common_all_to_all_embedding_trainable_v2(self, base_opt, test_opt, name):
               bp_v2=False,
               kv_creator=kv_creator,
               name='all2all_emb')
-          test_model = get_emb_sequential_model(tf.keras.layers.Embedding,
+          test_model = get_emb_sequential_model(layers.Embedding,
                                                 test_opt,
                                                 input_dim=start +
                                                 batch_size * i,
@@ -342,7 +351,7 @@ def common_all_to_all_embedding_trainable_v2(self, base_opt, test_opt, name):
         ckpt.save(save_dir + '/ckpt/test')
         del base_model
         del base_opt
-        tf.keras.backend.clear_session()
+        backend.clear_session()
         new_opt = de.DynamicEmbeddingOptimizer(Adam(1.1), synchronous=True)
         new_base_model = get_emb_sequential_model(
             de.keras.layers.HvdAllToAllEmbedding,
@@ -359,7 +368,7 @@ def common_all_to_all_embedding_trainable_v2(self, base_opt, test_opt, name):
         new_a2aemb_size = new_base_model.layers[0].params.size()
         self.assertEqual(a2aemb_size, new_a2aemb_size)
         hvd.join()  # Sync for avoiding files conflict
-        tf.keras.backend.clear_session()
+        backend.clear_session()
         new_base_model.load_weights(save_dir +
                                     '/variables/variables').expect_partial()
         new_a2aemb_size = new_base_model.layers[0].params.size()
@@ -390,7 +399,7 @@ def common_lazy_build_model_with_checkpoint_management_v2(self, name):
 
     dim = 8
 
-    class NoCompileModel(tf.keras.models.Model):
+    class NoCompileModel(models.Model):
 
       def __init__(self, init, dynamic=False):
         super().__init__(dynamic=dynamic)
@@ -400,8 +409,8 @@ def __init__(self, init, dynamic=False):
                                                         initializer=0,
                                                         kv_creator=kv_creator,
                                                         name=name)
-        self.l1 = tf.keras.layers.Dense(8, 'relu', kernel_initializer=init)
-        self.l2 = tf.keras.layers.Dense(1, 'sigmoid', kernel_initializer=init)
+        self.l1 = layers.Dense(8, 'relu', kernel_initializer=init)
+        self.l2 = layers.Dense(1, 'sigmoid', kernel_initializer=init)
 
       def build(self, input_shape):
         self.emb.build(input_shape)
@@ -499,7 +508,7 @@ def check_TFRADynamicEmbedding_directory(save_dir,
       del base_opt
       del model
       del ckpt
-      tf.keras.backend.clear_session()
+      backend.clear_session()
       tf.compat.v1.reset_default_graph()
 
       new_model = NoCompileModel('zeros')
@@ -558,7 +567,7 @@ def check_TFRADynamicEmbedding_directory(save_dir,
       del new_opt
       del new_model
       del new_ckpt
-      tf.keras.backend.clear_session()
+      backend.clear_session()
       tf.compat.v1.reset_default_graph()
       new_saved_model = NoCompileModel('zeros')
       new_saved_opt = Adam(1.2)
diff --git a/tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/shadow_embedding_ops_test.py b/tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/shadow_embedding_ops_test.py
index 4fecf226a..9acfe742a 100644
--- a/tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/shadow_embedding_ops_test.py
+++ b/tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/shadow_embedding_ops_test.py
@@ -48,9 +48,14 @@
 from tensorflow_recommenders_addons.utils.check_platform import is_macos, is_arm64
 
 try:
-  from tensorflow.keras.legacy.optimizers import Adam
+  from tf_keras import backend
+  from tf_keras.optimizers import Adam
 except:
-  from tensorflow.keras.optimizers import Adam
+  from tensorflow.keras import backend
+  try:
+    from tensorflow.keras.optimizers import Adam
+  except:
+    from tensorflow.keras.legacy.optimizers import Adam
 
 
 def _get_sparse_variable(name,
@@ -810,7 +815,7 @@ def size(self):
     shadow_value = module.shadow.read_value(False)
     self.assertAllEqual(shadow_value.shape, (0, 1))  # clear when saving
 
-    tf.keras.backend.clear_session()
+    backend.clear_session()
     del module, ckpt
     new_module = TestNewModule(table_devices_)
     new_ckpt = de.train.DECheckpoint(new_module)
@@ -821,7 +826,7 @@ def size(self):
     self.assertAllEqual(np.sort(expected_values, axis=0), values)
 
     # test expand shards number
-    tf.keras.backend.clear_session()
+    backend.clear_session()
     del new_module, new_ckpt
     shard_num = 5
     table_devices_ = table_device * shard_num
@@ -833,7 +838,7 @@ def size(self):
     self.assertAllEqual(np.sort(expected_values, axis=0), values)
 
     # test contracte shards number
-    tf.keras.backend.clear_session()
+    backend.clear_session()
     del new_module, new_ckpt
     shard_num = 2
     table_devices_ = table_device * shard_num
@@ -845,7 +850,7 @@ def size(self):
     self.assertAllEqual(np.sort(expected_values, axis=0), values)
 
     # test load all into one shard
-    tf.keras.backend.clear_session()
+    backend.clear_session()
     del new_module, new_ckpt
     shard_num = 1
     table_devices_ = table_device * shard_num
diff --git a/tensorflow_recommenders_addons/dynamic_embedding/python/ops/dynamic_embedding_optimizer.py b/tensorflow_recommenders_addons/dynamic_embedding/python/ops/dynamic_embedding_optimizer.py
index f3100000d..053bfc0df 100644
--- a/tensorflow_recommenders_addons/dynamic_embedding/python/ops/dynamic_embedding_optimizer.py
+++ b/tensorflow_recommenders_addons/dynamic_embedding/python/ops/dynamic_embedding_optimizer.py
@@ -63,12 +63,18 @@
   from tensorflow.python.training.tracking import base as trackable
 from tensorflow.python.keras.optimizer_v2 import optimizer_v2 as optimizer_v2_legacy
 from tensorflow.python.keras.optimizer_v2 import utils as optimizer_v2_legacy_utils
-try:  # Keras version >= 2.12.0
-  from tensorflow.keras.optimizers.legacy import Optimizer as keras_OptimizerV2_legacy
-  from tensorflow.keras.optimizers import Optimizer as keras_OptimizerV2
+
+try:  # tf version >= 2.16
+  from tf_keras.optimizers.legacy import Optimizer as keras_OptimizerV2_legacy
+  from tf_keras.optimizers import Optimizer as keras_OptimizerV2
 except:
-  from tensorflow.keras.optimizers import Optimizer as keras_OptimizerV2_legacy
-  keras_OptimizerV2 = keras_OptimizerV2_legacy
+  try:  # Keras version >= 2.12.0
+    from tensorflow.keras.optimizers.legacy import Optimizer as keras_OptimizerV2_legacy
+    from tensorflow.keras.optimizers import Optimizer as keras_OptimizerV2
+  except:
+    from tensorflow.keras.optimizers import Optimizer as keras_OptimizerV2_legacy
+    keras_OptimizerV2 = keras_OptimizerV2_legacy
+
 from tensorflow.python.eager import tape
 from tensorflow.python.distribute import values_util as distribute_values_util
 from tensorflow.python.distribute import distribute_utils
diff --git a/tensorflow_recommenders_addons/dynamic_embedding/python/ops/dynamic_embedding_variable.py b/tensorflow_recommenders_addons/dynamic_embedding/python/ops/dynamic_embedding_variable.py
index 4d5d7e339..7e411c9a7 100644
--- a/tensorflow_recommenders_addons/dynamic_embedding/python/ops/dynamic_embedding_variable.py
+++ b/tensorflow_recommenders_addons/dynamic_embedding/python/ops/dynamic_embedding_variable.py
@@ -55,7 +55,22 @@
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
+
 from tensorflow.python.keras.optimizer_v2.optimizer_v2 import OptimizerV2
+
+try:  # tf version >= 2.16
+  from tf_keras.initializers import Initializer
+  from tf_keras.optimizers.legacy import Optimizer as keras_OptimizerV2_legacy
+  from tf_keras.optimizers import Optimizer as keras_OptimizerV2
+except:
+  from tensorflow.keras.initializers import Initializer
+  try:  # Keras version >= 2.12.0
+    from tensorflow.keras.optimizers.legacy import Optimizer as keras_OptimizerV2_legacy
+    from tensorflow.keras.optimizers import Optimizer as keras_OptimizerV2
+  except:
+    from tensorflow.keras.optimizers import Optimizer as keras_OptimizerV2_legacy
+    keras_OptimizerV2 = keras_OptimizerV2_legacy
+
 from tensorflow.python.keras.utils import tf_utils
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import bitwise_ops
@@ -696,10 +711,7 @@ def restrict_policy(self):
 
   def _convert_anything_to_init(self, raw_init, dim):
     init = raw_init
-    valid_list = [
-        init_ops.Initializer, init_ops_v2.Initializer,
-        tf.keras.initializers.Initializer
-    ]
+    valid_list = [init_ops.Initializer, init_ops_v2.Initializer, Initializer]
     if kinit2 is not None:
       valid_list.append(kinit2.Initializer)
     valid_list = tuple(valid_list)
@@ -1145,8 +1157,7 @@ def get_slot_variables(self, optimizer):
     Returns:
       List of slot `Variable`s in optimizer.
     """
-    if not isinstance(optimizer,
-                      (Optimizer, OptimizerV2, tf.keras.optimizers.Optimizer)):
+    if not isinstance(optimizer, (Optimizer, OptimizerV2, keras_OptimizerV2)):
       raise TypeError('Expect an optimizer, but get {}'.format(type(optimizer)))
     slots = []
     if hasattr(optimizer, 'get_slot_names'):
diff --git a/tensorflow_recommenders_addons/dynamic_embedding/python/ops/tf_patch.py b/tensorflow_recommenders_addons/dynamic_embedding/python/ops/tf_patch.py
index 102e484e3..a1dcb8bf7 100644
--- a/tensorflow_recommenders_addons/dynamic_embedding/python/ops/tf_patch.py
+++ b/tensorflow_recommenders_addons/dynamic_embedding/python/ops/tf_patch.py
@@ -14,7 +14,6 @@
 
 # lint-as: python3
 """patch on tensorflow"""
-
 from tensorflow_recommenders_addons import dynamic_embedding as de
 
 try:
@@ -24,11 +23,14 @@
   pass  # for compatible with TF < 2.3.x
 
 try:
-  import tensorflow as tf
-  kinit_tf = tf.keras.initializers
-except ImportError:
-  kinit_tf = None
-  pass  # for compatible with TF >= 2.6.x
+  import tf_keras
+  kinit_tf = tf_keras.initializers
+except:
+  try:
+    import tensorflow
+    kinit_tf = tensorflow.keras.initializers
+  except ImportError:
+    pass  # for compatible with TF >= 2.6.x
 
 try:
   import keras as tmp_keras
diff --git a/tensorflow_recommenders_addons/utils/test_utils.py b/tensorflow_recommenders_addons/utils/test_utils.py
index 2394e1b44..81bcebf01 100644
--- a/tensorflow_recommenders_addons/utils/test_utils.py
+++ b/tensorflow_recommenders_addons/utils/test_utils.py
@@ -28,6 +28,11 @@
   from distutils.version import LooseVersion as Version
 from tensorflow_recommenders_addons.utils import resource_loader
 
+try:
+  from tf_keras import mixed_precision
+except:
+  from tensorflow.keras import mixed_precision
+
 NUMBER_OF_WORKERS = int(os.environ.get("PYTEST_XDIST_WORKER_COUNT", "1"))
 WORKER_ID = int(os.environ.get("PYTEST_XDIST_WORKER", "gw0")[2])
 NUMBER_OF_GPUS = len(resource_loader.get_devices("GPU"))
@@ -106,9 +111,9 @@ def only_run_functions_eagerly(request):
 def run_with_mixed_precision_policy(request):
   if is_gpu_available() and Version(tf.__version__) <= "2.4.1":
     pytest.xfail("See https://github.com/tensorflow/tensorflow/issues/39775")
-  tf.keras.mixed_precision.experimental.set_policy(request.param)
+  mixed_precision.experimental.set_policy(request.param)
   yield
-  tf.keras.mixed_precision.experimental.set_policy("float32")
+  mixed_precision.experimental.set_policy("float32")
 
 
 @pytest.fixture(scope="function", params=["channels_first", "channels_last"])
diff --git a/tensorflow_recommenders_addons/utils/tests/test_utils_test.py b/tensorflow_recommenders_addons/utils/tests/test_utils_test.py
index 4476a00aa..c3b33a938 100644
--- a/tensorflow_recommenders_addons/utils/tests/test_utils_test.py
+++ b/tensorflow_recommenders_addons/utils/tests/test_utils_test.py
@@ -9,6 +9,11 @@
 import tensorflow as tf
 from tensorflow_recommenders_addons.utils import test_utils
 
+try:
+  from tf_keras import layers, Model
+except:
+  from tensorflow.keras import layers, Model
+
 
 def test_seed_is_set():
   assert random.randint(0, 10000) == 6311
@@ -24,9 +29,9 @@ def test_all_scopes(device):
 
 
 def train_small_model():
-  model_input = tf.keras.layers.Input((3,))
-  model_output = tf.keras.layers.Dense(4)(model_input)
-  model = tf.keras.Model(model_input, model_output)
+  model_input = layers.Input((3,))
+  model_output = layers.Dense(4)(model_input)
+  model = Model(model_input, model_output)
   model.compile(loss="mse")
 
   x = np.random.uniform(size=(5, 3))
diff --git a/tensorflow_recommenders_addons/utils/types.py b/tensorflow_recommenders_addons/utils/types.py
index d30a2d72b..bb9c3508c 100644
--- a/tensorflow_recommenders_addons/utils/types.py
+++ b/tensorflow_recommenders_addons/utils/types.py
@@ -20,6 +20,12 @@
 import numpy as np
 import tensorflow as tf
 
+try:  # tf version >= 2.16
+  from tf_keras.optimizers import Optimizer as keras_OptimizerV2
+except:
+  # Keras version >= 2.12.0
+  from tensorflow.keras.optimizers import Optimizer as keras_OptimizerV2
+
 Number = Union[
     float,
     int,
@@ -40,7 +46,7 @@
 Regularizer = Union[None, dict, str, Callable]
 Constraint = Union[None, dict, str, Callable]
 Activation = Union[None, str, Callable]
-Optimizer = Union[tf.keras.optimizers.Optimizer, str]
+Optimizer = Union[keras_OptimizerV2, str]
 
 TensorLike = Union[
     List[Union[Number, list]],
diff --git a/tools/docker/build_wheel.Dockerfile b/tools/docker/build_wheel.Dockerfile
index 5f8a8e4ab..8ba7a113e 100644
--- a/tools/docker/build_wheel.Dockerfile
+++ b/tools/docker/build_wheel.Dockerfile
@@ -53,6 +53,10 @@ RUN python -m pip install -r requirements.txt
 
 RUN python -m pip install tensorflow-io
 
+RUN bash -c 'if [[ "$TF_VERSION" =~ ^2\.(16|17)\.[0-9]+$ ]]; then \
+          pip install tf-keras==2.16.0; \
+      fi'
+
 RUN python -m pip install --upgrade protobuf==$PROTOBUF_VERSION
 
 COPY ./ /recommenders-addons
diff --git a/tools/docker/cuda12.3-cudnn8.9-ubuntu20.04-manylinux2014-python3.10.Dockerfile b/tools/docker/cuda12.3-cudnn8.9-ubuntu20.04-manylinux2014-python3.10.Dockerfile
index 010320634..a29b3d0e6 100644
--- a/tools/docker/cuda12.3-cudnn8.9-ubuntu20.04-manylinux2014-python3.10.Dockerfile
+++ b/tools/docker/cuda12.3-cudnn8.9-ubuntu20.04-manylinux2014-python3.10.Dockerfile
@@ -6,8 +6,9 @@
 # To push a new version, run:
 # $ docker build -f cuda12.3-cudnn8.9-ubuntu20.04-manylinux2014-python3.10.Dockerfile . \
 #  --tag "tfra/nosla-cuda12.3-cudnn8.9-ubuntu20.04-manylinux2014-python3.10"
+# docker login -u user -p token
 # $ docker push tfra/nosla-cuda12.3-cudnn8.9-ubuntu20.04-manylinux2014-python3.10
-FROM nvidia/cuda:12.3.2-devel-ubuntu20.04 as devtoolset
+FROM docker.io/tfra/nosla-cuda12.3.2-cudnn8.9-devel-ubuntu20.04 AS devtoolset
 
 RUN chmod 777 /tmp/
 ENV DEBIAN_FRONTEND=noninteractive
@@ -41,7 +42,7 @@ ADD devtoolset/rpm-patch.sh rpm-patch.sh
 RUN /build_devtoolset.sh devtoolset-8 /dt8
 
 # TODO(klimek): Split up into two different docker images.
-FROM nvidia/cuda:12.3.2-devel-ubuntu20.04
+FROM docker.io/tfra/nosla-cuda12.3.2-cudnn8.9-devel-ubuntu20.04
 COPY --from=devtoolset /dt8 /dt8
 
 # Install TensorRT.
diff --git a/tools/docker/cuda12.3-cudnn8.9-ubuntu20.04-manylinux2014-python3.11.Dockerfile b/tools/docker/cuda12.3-cudnn8.9-ubuntu20.04-manylinux2014-python3.11.Dockerfile
index 9e275e8b5..243e43e14 100644
--- a/tools/docker/cuda12.3-cudnn8.9-ubuntu20.04-manylinux2014-python3.11.Dockerfile
+++ b/tools/docker/cuda12.3-cudnn8.9-ubuntu20.04-manylinux2014-python3.11.Dockerfile
@@ -6,8 +6,9 @@
 # To push a new version, run:
 # $ docker build -f cuda12.3-cudnn8.9-ubuntu20.04-manylinux2014-python3.11.Dockerfile . \
 #  --tag "tfra/nosla-cuda12.3-cudnn8.9-ubuntu20.04-manylinux2014-python3.11"
+# docker login -u user -p token
 # $ docker push tfra/nosla-cuda12.3-cudnn8.9-ubuntu20.04-manylinux2014-python3.11
-FROM nvidia/cuda:12.3.2-devel-ubuntu20.04 as devtoolset
+FROM docker.io/tfra/nosla-cuda12.3.2-cudnn8.9-devel-ubuntu20.04 AS devtoolset
 
 RUN chmod 777 /tmp/
 ENV DEBIAN_FRONTEND=noninteractive
@@ -41,7 +42,7 @@ ADD devtoolset/rpm-patch.sh rpm-patch.sh
 RUN /build_devtoolset.sh devtoolset-8 /dt8
 
 # TODO(klimek): Split up into two different docker images.
-FROM nvidia/cuda:12.3.2-devel-ubuntu20.04
+FROM docker.io/tfra/nosla-cuda12.3.2-cudnn8.9-devel-ubuntu20.04
 COPY --from=devtoolset /dt8 /dt8
 
 # Install TensorRT.
diff --git a/tools/docker/cuda12.3-cudnn8.9-ubuntu20.04-manylinux2014-python3.9.Dockerfile b/tools/docker/cuda12.3-cudnn8.9-ubuntu20.04-manylinux2014-python3.9.Dockerfile
index 6f32e4b09..fde2b931a 100644
--- a/tools/docker/cuda12.3-cudnn8.9-ubuntu20.04-manylinux2014-python3.9.Dockerfile
+++ b/tools/docker/cuda12.3-cudnn8.9-ubuntu20.04-manylinux2014-python3.9.Dockerfile
@@ -6,8 +6,9 @@
 # To push a new version, run:
 # $ docker build -f cuda12.3-cudnn8.9-ubuntu20.04-manylinux2014-python3.9.Dockerfile . \
 #  --tag "tfra/nosla-cuda12.3-cudnn8.9-ubuntu20.04-manylinux2014-python3.9"
-# $ docker push tfra/nosla-12.3-cudnn8.9-ubuntu20.04-manylinux2014-python3.9
-FROM nvidia/cuda:12.3.2-devel-ubuntu20.04 as devtoolset
+# docker login -u user -p token
+# $ docker push tfra/nosla-cuda12.3-cudnn8.9-ubuntu20.04-manylinux2014-python3.9
+FROM docker.io/tfra/nosla-cuda12.3.2-cudnn8.9-devel-ubuntu20.04 AS devtoolset
 
 RUN chmod 777 /tmp/
 ENV DEBIAN_FRONTEND=noninteractive
@@ -41,7 +42,7 @@ ADD devtoolset/rpm-patch.sh rpm-patch.sh
 RUN /build_devtoolset.sh devtoolset-8 /dt8
 
 # TODO(klimek): Split up into two different docker images.
-FROM nvidia/cuda:12.2.2-cudnn8-devel-ubuntu20.04
+FROM docker.io/tfra/nosla-cuda12.3.2-cudnn8.9-devel-ubuntu20.04
 COPY --from=devtoolset /dt8 /dt8
 
 # Install TensorRT.
diff --git a/tools/docker/cuda12.3.2-cudnn8.9-devel-ubuntu20.04.Dockerfile b/tools/docker/cuda12.3.2-cudnn8.9-devel-ubuntu20.04.Dockerfile
index 81121378b..1cd91740a 100644
--- a/tools/docker/cuda12.3.2-cudnn8.9-devel-ubuntu20.04.Dockerfile
+++ b/tools/docker/cuda12.3.2-cudnn8.9-devel-ubuntu20.04.Dockerfile
@@ -1,28 +1,18 @@
 # To push a new version, run:
 # $ docker build -f cuda12.3.2-cudnn8.9-devel-ubuntu20.04.Dockerfile . --tag "tfra/nosla-cuda12.3.2-cudnn8.9-devel-ubuntu20.04"
-# $ docker push tfra/nosla-12.3.2-cudnn8.9-devel-ubuntu20.04
-
-# downlod https://developer.nvidia.com/downloads/compute/cudnn/secure/8.9.7/local_installers/12.x/cudnn-local-repo-ubuntu2204-8.9.7.29_1.0-1_amd64.deb/
-# en
-FROM nvidia/cuda:12.3.2-devel-ubuntu20.04 as devtoolset
+# $ docker login -u user -p token
+# $ docker push tfra/nosla-cuda12.3.2-cudnn8.9-devel-ubuntu20.04
+# https://developer.nvidia.com/rdp/cudnn-archive to find correct cudnn.
+# https://developer.nvidia.com/downloads/compute/cudnn/secure/8.9.7/local_installers/12.x/cudnn-local-repo-ubuntu2004-8.9.7.29_1.0-1_amd64.deb/
+# dowloand to local and copy to the docker build context.
+FROM nvidia/cuda:12.3.2-devel-ubuntu20.04 AS devtoolset
 RUN apt-get update
-#RUN apt-get update && apt-cache search linux-headers
-
-#RUN apt-get install linux-headers-$(uname -r)
-#RUN apt-key del 7fa2af80
-RUN rm -rf /usr/share/keyrings/cuda-archive-keyring.gpg
 RUN apt-get -y install wget
-RUN wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-keyring_1.1-1_all.deb
-
-RUN dpkg -i cuda-keyring_1.1-1_all.deb  && rm cuda-keyring_1.1-1_all.deb
-RUN cat /etc/apt/sources.list && \
-    ls -l /etc/apt/sources.list.d/
+COPY cudnn-local-repo-ubuntu2004-8.9.7.29_1.0-1_amd64.deb /tmp/
 
-RUN rm -f /etc/apt/sources.list.d/cuda*
+RUN dpkg -i /tmp/cudnn-local-repo-ubuntu2004-8.9.7.29_1.0-1_amd64.deb
 
-RUN echo "deb [signed-by=/usr/share/keyrings/cuda-archive-keyring.gpg] https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/ /" > /etc/apt/sources.list.d/cuda.list
-
-RUN apt-get update && \
-    apt-get -y install cudnn8-cuda-12
-
-RUN apt-get clean && rm -rf /var/lib/apt/lists/* /tmp/*
\ No newline at end of file
+RUN cp /var/cudnn-local-repo-ubuntu2004-8.9.7.29/cudnn-*-keyring.gpg /usr/share/keyrings/
+RUN apt-get update
+RUN apt-get -y install /var/cudnn-local-repo-ubuntu2004-8.9.7.29/libcudnn*.deb
+RUN apt-get clean && rm -rf /var/lib/apt/lists/* /tmp/* /var/cudnn-local-repo-ubuntu2004-8.9.7.29/*.deb
\ No newline at end of file
diff --git a/tools/install_deps/tf-keras.txt b/tools/install_deps/tf-keras.txt
deleted file mode 100644
index c9b91775c..000000000
--- a/tools/install_deps/tf-keras.txt
+++ /dev/null
@@ -1,3 +0,0 @@
-tensorflow>=2.16.0,<=2.16.2
-tensorflow-estimator
-tf-keras<=2.16.0
\ No newline at end of file