|
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