Skip to content

Commit 08f9b1a

Browse files
qlzh727tensorflower-gardener
authored andcommitted
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
1 parent ed008ff commit 08f9b1a

File tree

3 files changed

+1
-10
lines changed

3 files changed

+1
-10
lines changed

keras/api/golden/v2/tensorflow.keras.dtensor.experimental.pbtxt

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,4 @@ tf_module {
88
name: "optimizers"
99
mtype: "<type \'module\'>"
1010
}
11-
member_method {
12-
name: "layout_map_scope"
13-
argspec: "args=[], varargs=args, keywords=kwds, defaults=None"
14-
}
1511
}

keras/dtensor/layout_map.py

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,6 @@
2727
from keras.engine import base_layer
2828

2929
# isort: off
30-
from tensorflow.python.util.deprecation import deprecated
3130
from tensorflow.python.util.tf_export import keras_export
3231

3332

@@ -235,10 +234,6 @@ def call(self, inputs):
235234
LayoutMap.get.__doc__ = LayoutMap.__getitem__.__doc__
236235

237236

238-
@keras_export("keras.dtensor.experimental.layout_map_scope", v1=[])
239-
@deprecated(
240-
None, "use tf.keras.dtensor.experimental.LayoutMap.scope() instead."
241-
)
242237
@contextlib.contextmanager
243238
def layout_map_scope(layout_map):
244239
"""Apply the layout to all the tf.Variables created under the scope.

keras/dtensor/layout_map_test.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -340,7 +340,7 @@ def test_init_model_with_empty_layout_map(self):
340340

341341
def test_weight_regularization(self):
342342
layout_map = layout_map_lib.LayoutMap(mesh=self.mesh)
343-
with layout_map_lib.layout_map_scope(layout_map):
343+
with layout_map.scope():
344344
model = models.Sequential(
345345
[
346346
layers.Dense(

0 commit comments

Comments
 (0)