Skip to content

Commit 3ad57a2

Browse files
qlzh727copybara-github
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 2916873 commit 3ad57a2

File tree

1 file changed

+4
-4
lines changed

1 file changed

+4
-4
lines changed

site/en/tutorials/distribute/dtensor_keras_tutorial.ipynb

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -634,7 +634,7 @@
634634
"layout_map['feature.*kernel'] = dtensor.Layout.batch_sharded(mesh, 'batch', rank=2)\n",
635635
"layout_map['feature.*bias'] = dtensor.Layout.batch_sharded(mesh, 'batch', rank=1)\n",
636636
"\n",
637-
"with tf.keras.dtensor.experimental.layout_map_scope(layout_map):\n",
637+
"with layout_map.scope():\n",
638638
" subclassed_model = SubclassedModel()"
639639
]
640640
},
@@ -713,7 +713,7 @@
713713
},
714714
"outputs": [],
715715
"source": [
716-
"with tf.keras.dtensor.experimental.layout_map_scope(layout_map):\n",
716+
"with layout_map.scope():\n",
717717
" inputs = tf.keras.Input((16,), batch_size=16)\n",
718718
" x = tf.keras.layers.Dense(16, name='feature')(inputs)\n",
719719
" x = tf.keras.layers.Dropout(0.1)(x)\n",
@@ -731,7 +731,7 @@
731731
},
732732
"outputs": [],
733733
"source": [
734-
"with tf.keras.dtensor.experimental.layout_map_scope(layout_map):\n",
734+
"with layout_map.scope():\n",
735735
" model = tf.keras.Sequential([\n",
736736
" tf.keras.layers.Dense(16, name='feature', input_shape=(16,)),\n",
737737
" tf.keras.layers.Dropout(0.1),\n",
@@ -744,8 +744,8 @@
744744
],
745745
"metadata": {
746746
"colab": {
747-
"collapsed_sections": [],
748747
"name": "dtensor_keras_tutorial.ipynb",
748+
"provenance": [],
749749
"toc_visible": true
750750
},
751751
"kernelspec": {

0 commit comments

Comments
 (0)