From 08f9b1a2f8248b2158b50d02312129551997f596 Mon Sep 17 00:00:00 2001
From: Scott Zhu <scottzhu@google.com>
Date: Mon, 24 Apr 2023 19:18:37 -0700
Subject: [PATCH] Remove the deprecated
 `tf.keras.dtensor.experimental.layout_map_scope` API.

The deprecated warning was showing up when using the `LayoutMap.scope`, which is confusing to end user.

PiperOrigin-RevId: 526818762
---
 .../golden/v2/tensorflow.keras.dtensor.experimental.pbtxt    | 4 ----
 keras/dtensor/layout_map.py                                  | 5 -----
 keras/dtensor/layout_map_test.py                             | 2 +-
 3 files changed, 1 insertion(+), 10 deletions(-)

diff --git a/keras/api/golden/v2/tensorflow.keras.dtensor.experimental.pbtxt b/keras/api/golden/v2/tensorflow.keras.dtensor.experimental.pbtxt
index 20f3bd29b56..dd963f6657d 100644
--- a/keras/api/golden/v2/tensorflow.keras.dtensor.experimental.pbtxt
+++ b/keras/api/golden/v2/tensorflow.keras.dtensor.experimental.pbtxt
@@ -8,8 +8,4 @@ tf_module {
     name: "optimizers"
     mtype: "<type \'module\'>"
   }
-  member_method {
-    name: "layout_map_scope"
-    argspec: "args=[], varargs=args, keywords=kwds, defaults=None"
-  }
 }
diff --git a/keras/dtensor/layout_map.py b/keras/dtensor/layout_map.py
index 49476c00f2a..c7fd3407d53 100644
--- a/keras/dtensor/layout_map.py
+++ b/keras/dtensor/layout_map.py
@@ -27,7 +27,6 @@
 from keras.engine import base_layer
 
 # isort: off
-from tensorflow.python.util.deprecation import deprecated
 from tensorflow.python.util.tf_export import keras_export
 
 
@@ -235,10 +234,6 @@ def call(self, inputs):
 LayoutMap.get.__doc__ = LayoutMap.__getitem__.__doc__
 
 
-@keras_export("keras.dtensor.experimental.layout_map_scope", v1=[])
-@deprecated(
-    None, "use tf.keras.dtensor.experimental.LayoutMap.scope() instead."
-)
 @contextlib.contextmanager
 def layout_map_scope(layout_map):
     """Apply the layout to all the tf.Variables created under the scope.
diff --git a/keras/dtensor/layout_map_test.py b/keras/dtensor/layout_map_test.py
index 268180a14ce..7df61a78d47 100644
--- a/keras/dtensor/layout_map_test.py
+++ b/keras/dtensor/layout_map_test.py
@@ -340,7 +340,7 @@ def test_init_model_with_empty_layout_map(self):
 
     def test_weight_regularization(self):
         layout_map = layout_map_lib.LayoutMap(mesh=self.mesh)
-        with layout_map_lib.layout_map_scope(layout_map):
+        with layout_map.scope():
             model = models.Sequential(
                 [
                     layers.Dense(