|
634 | 634 | "layout_map['feature.*kernel'] = dtensor.Layout.batch_sharded(mesh, 'batch', rank=2)\n",
|
635 | 635 | "layout_map['feature.*bias'] = dtensor.Layout.batch_sharded(mesh, 'batch', rank=1)\n",
|
636 | 636 | "\n",
|
637 |
| - "with tf.keras.dtensor.experimental.layout_map_scope(layout_map):\n", |
| 637 | + "with layout_map.scope():\n", |
638 | 638 | " subclassed_model = SubclassedModel()"
|
639 | 639 | ]
|
640 | 640 | },
|
|
713 | 713 | },
|
714 | 714 | "outputs": [],
|
715 | 715 | "source": [
|
716 |
| - "with tf.keras.dtensor.experimental.layout_map_scope(layout_map):\n", |
| 716 | + "with layout_map.scope():\n", |
717 | 717 | " inputs = tf.keras.Input((16,), batch_size=16)\n",
|
718 | 718 | " x = tf.keras.layers.Dense(16, name='feature')(inputs)\n",
|
719 | 719 | " x = tf.keras.layers.Dropout(0.1)(x)\n",
|
|
731 | 731 | },
|
732 | 732 | "outputs": [],
|
733 | 733 | "source": [
|
734 |
| - "with tf.keras.dtensor.experimental.layout_map_scope(layout_map):\n", |
| 734 | + "with layout_map.scope():\n", |
735 | 735 | " model = tf.keras.Sequential([\n",
|
736 | 736 | " tf.keras.layers.Dense(16, name='feature', input_shape=(16,)),\n",
|
737 | 737 | " tf.keras.layers.Dropout(0.1),\n",
|
|
744 | 744 | ],
|
745 | 745 | "metadata": {
|
746 | 746 | "colab": {
|
747 |
| - "collapsed_sections": [], |
748 | 747 | "name": "dtensor_keras_tutorial.ipynb",
|
| 748 | + "provenance": [], |
749 | 749 | "toc_visible": true
|
750 | 750 | },
|
751 | 751 | "kernelspec": {
|
|
0 commit comments