site/ko/guide/dtensor_overview.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1ljvLya59ep5"
      },
      "source": [
        "##### Copyright 2019 The TensorFlow Authors.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tuOe1ymfHZPu"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VcQIa1uG86Wh"
      },
      "source": [
        "# DTensor 개념"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6dWNQEum9AfY"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>     <a target=\"_blank\" href=\"https://www.tensorflow.org/guide/dtensor_overview\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\"> TensorFlow.org에서 보기</a> </td>\n",
        "  <td>     <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs-l10n/blob/master/site/ko/guide/dtensor_overview.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\">Google Colab에서 실행하기</a> </td>\n",
        "  <td><a target=\"_blank\" href=\"https://github.com/tensorflow/docs-l10n/blob/master/site/ko/guide/dtensor_overview.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\">GitHub에서 소스 보기</a></td>\n",
        "  <td><a href=\"https://storage.googleapis.com/tensorflow_docs/docs-l10n/site/ko/guide/dtensor_overview.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\">노트북 다운로드하기</a></td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MGZuakHVlVQf"
      },
      "source": [
        "## 개요\n",
        "\n",
        "이 Colab에서는 동기식 분산 컴퓨팅을 지원하며 TensorFlow의 확장 버전인 DTensor를 소개합니다.\n",
        "\n",
        "DTensor는 개발자가 내부에서 장치 전반의 배포를 관리하면서 글로벌적으로 텐서에서 작동하는 애플리케이션을 구성할 수 있도록 하는 글로벌 프로그래밍 모델을 제공합니다. DTensor는 분할 지시문에 따라 *[단일 프로그램 다중데이터(SPMD)](https://en.wikipedia.org/wiki/SPMD) 확장*이라는 절차를 통해 프로그램과 텐서를 배포합니다\n",
        "\n",
        "분할 지시문으로부터 애플리케이션을 분리함으로써 DTensor는 글로벌 의미 체계를 유지하면서도 단일 기기, 여러 기기 또는 여러 클라이언트에서 동일한 애플리케이션을 실행할 수 있습니다.\n",
        "\n",
        "이 가이드에서는 분산 컴퓨팅을 위한 DTensor 개념과 DTensor가 TensorFlow와 통합되는 방식을 소개합니다. 모델 훈련에서 DTensor를 사용하는 데모를 확인하려면 [DTensor를 사용한 분산 훈련](https://www.tensorflow.org/tutorials/distribute/dtensor_ml_tutorial) 튜토리얼을 참조합니다."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h7ZTDq7KngwA"
      },
      "source": [
        "## 설치하기\n",
        "\n",
        "DTensor는 TensorFlow 2.9.0 릴리스의 일부이며 2022년 4월 9일부터 TensorFlow 나이틀리 빌드에도 포함되어 있습니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OKaPw8vwwZAC"
      },
      "outputs": [],
      "source": [
        "!pip install --quiet --upgrade --pre tensorflow"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O3pG29uZIWYO"
      },
      "source": [
        "설치가 완료되면 `tensorflow` 및 `tf.experimental.dtensor`를 가져옵니다. 그런 다음 6개의 가상 CPU를 사용하도록 TensorFlow를 구성합니다.\n",
        "\n",
        "이 예제에서는 vCPU를 사용하지만 DTensor는 CPU, GPU 또는 TPU 장치에서 동일한 방식으로 작동합니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Q92lo0zjwej8"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "from tensorflow.experimental import dtensor\n",
        "\n",
        "print('TensorFlow version:', tf.__version__)\n",
        "\n",
        "def configure_virtual_cpus(ncpu):\n",
        "  phy_devices = tf.config.list_physical_devices('CPU')\n",
        "  tf.config.set_logical_device_configuration(phy_devices[0], [\n",
        "        tf.config.LogicalDeviceConfiguration(),\n",
        "    ] * ncpu)\n",
        "\n",
        "configure_virtual_cpus(6)\n",
        "DEVICES = [f'CPU:{i}' for i in range(6)]\n",
        "\n",
        "tf.config.list_logical_devices('CPU')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O-lsrxUnlsCC"
      },
      "source": [
        "## DTensor의 분산 텐서 모델\n",
        "\n",
        "DTensor는 `dtensor.Mesh` 및 `dtensor.Layout`이라는 두 가지 개념을 소개합니다. 이는 토폴로지와 관련된 장치에서 텐서의 분할을 모델링하는 추상화입니다.\n",
        "\n",
        "- `Mesh`는 계산용 기기 목록을 정의합니다.\n",
        "- `Layout`은 `Mesh`의 텐서 차원을 분할하는 방식을 정의합니다."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JjiHaH0ql9yo"
      },
      "source": [
        "### 메시(Mesh)\n",
        "\n",
        "`Mesh`는 장치 세트의 논리적 데카르트 토폴로지를 나타냅니다. 데카르트 그리드의 각 차원은 **메시 차원**이라고 하며 이름으로 참조합니다. 동일한 `Mesh` 내 메시 차원의 이름은 고유해야 합니다.\n",
        "\n",
        "메시 차원의 이름은 각 축을 따라 `tf.Tensor`의 분할 동작을 설명할 수 있도록 `Layout`을 통해 참조됩니다. 이는 나중에 `Layout` 섹션에서 자세히 설명합니다.\n",
        "\n",
        "`Mesh`는 장치의 다차원 배열로 생각할 수도 있습니다."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_J6cOieEbaUw"
      },
      "source": [
        "1차원 `Mesh`에서는 모든 장치가 단일 메시 차원으로 목록을 형성합니다. 다음 예제에서는 6개 장치를 사용하는 `'x'` 메시 차원에 따라 `dtensor.create_mesh`를 사용하여 메시를 생성합니다.\n",
        "\n",
        "\n",
        "<img src=\"https://www.tensorflow.org/images/dtensor/dtensor_mesh_1d.png\" class=\"no-filter\" alt=\"6개의 CPU를 사용하는 1차원 메시\">      \n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QLH5fgdBmA58"
      },
      "outputs": [],
      "source": [
        "mesh_1d = dtensor.create_mesh([('x', 6)], devices=DEVICES)\n",
        "print(mesh_1d)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hSZwaUwnEgXB"
      },
      "source": [
        "`Mesh`는 다차원일 수도 있습니다. 다음 예제에서 6개의 CPU 장치는 `3x2` 메시를 형성합니다. 여기에서 `'x'` 메시 차원은 3개 장치의 크기를 가지며 `'y'` 메시 차원은 2개 장치의 크기를 갖습니다.\n",
        "\n",
        "<img src=\"https://www.tensorflow.org/images/dtensor/dtensor_mesh_2d.png\" alt=\"A 2 dimensional mesh with 6 CPUs\" class=\"\">"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "op6TmKUQE-sZ"
      },
      "outputs": [],
      "source": [
        "mesh_2d = dtensor.create_mesh([('x', 3), ('y', 2)], devices=DEVICES)\n",
        "print(mesh_2d)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "deAqdrDPFn2f"
      },
      "source": [
        "### 레이아웃(Layout)\n",
        "\n",
        "**`Layout`**은 텐서가 `Mesh`에서 분산 또는 분할되는 방식을 지정합니다.\n",
        "\n",
        "참고: `Mesh`와 `Layout` 사이의 혼란을 피하기 위해 이 가이드에서 *차원(dimension)*이라는 용어는 항상 `Mesh`와 연결하여 사용하고 *축(axis)*라는 용어는 `Tensor` 및 `Layout`과 함께 사용합니다.\n",
        "\n",
        "`Layout`의 순위는 `Layout`이 적용된 `Tensor`의 순위와 같아야 합니다. `Tensor`의 각 축에 대해 `Layout`은 텐서를 분할할 메시 차원을 지정하거나 축을 '분할하지 않음'으로 지정할 수 있습니다. 텐서는 분할되지 않은 모든 메시 차원에 복제됩니다.\n",
        "\n",
        "`Layout`의 순위와 `Mesh`의 차원 수는 일치하지 않아도 됩니다. `Layout`의 `unsharded` 축은 메시 차원과 연결할 필요가 없으며 `unsharded` 메시 차원은 `layout` 축과 연결할 필요가 없습니다.\n",
        "\n",
        "<img src=\"https://www.tensorflow.org/images/dtensor/dtensor_components_diag.png\" alt=\"Diagram of dtensor components.\" class=\"\">"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Px_bF1c-bQ7e"
      },
      "source": [
        "이전 섹션에서 생성한 `Mesh`에 대한 `Layout`의 몇 가지 예제를 분석해 보겠습니다."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fqzCNlWAbm-c"
      },
      "source": [
        "`[(\"x\", 6)]`와 같은 1차원 메시(이전 섹션의 `mesh_1d`)에서 `Layout([\"unsharded\", \"unsharded\"], mesh_1d)`는 6개의 장치에 복제된 2순위 텐서의 레이아웃입니다. <img src=\"https://www.tensorflow.org/images/dtensor/dtensor_layout_replicated.png\" class=\"no-filter\" alt=\"1순위 메시에 복제된 텐서\">"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-a3EnmZag6x1"
      },
      "outputs": [],
      "source": [
        "layout = dtensor.Layout([dtensor.UNSHARDED, dtensor.UNSHARDED], mesh_1d)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ywRJwuLDt2Qq"
      },
      "source": [
        "동일한 텐서를 사용하고 `Layout(['unsharded', 'x'])` 레이아웃을 메시하면 6개의 장치에서 텐서의 두 번째 축이 분할될 수 있습니다.\n",
        "\n",
        "\n",
        "<img src=\"https://www.tensorflow.org/images/dtensor/dtensor_layout_rank1.png\" class=\"no-filter\" alt=\"1순위 메시에서 분할된 텐서\">      "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7BgqL0jUvV5a"
      },
      "outputs": [],
      "source": [
        "layout = dtensor.Layout([dtensor.UNSHARDED, 'x'], mesh_1d)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DgciDNmK76l9"
      },
      "source": [
        "`[(\"x\", 3), (\"y\", 2)]`(이전 섹션의 `mesh_2d`)와 같은 2차원의 3x2 메시가 주어질 경우 { code2}Layout([\"y\", \"x\"], mesh_2d)는 첫 번째 축이 메시 차원 `\"y\"`에 걸쳐 샤딩되는 2순위 `Tensor`의 레이아웃이며 두 번째 축은 메시 차원 `\"x\"`에 걸쳐 샤딩됩니다."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Eyp_qOSyvieo"
      },
      "source": [
        "<img src=\"https://www.tensorflow.org/images/dtensor/dtensor_layout_rank2.png\" class=\"no-filter\" alt=\"메시 차원 'y'에서 분할된 첫 번째 축과 메시 차원 'x'에서 분할된 두 번째 축이 있는 텐서\">      \n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "p8OrehEuhPbS"
      },
      "outputs": [],
      "source": [
        "layout = dtensor.Layout(['y', 'x'], mesh_2d)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1Kyg0V3ehMNJ"
      },
      "source": [
        "동일한 `mesh_2d`의 경우 레이아웃 `Layout([\"x\", dtensor.UNSHARDED], mesh_2d)`은 2순위 `Tensor`이며, 이는 `\"y\"`에 복제되고 첫 번째 축이 메시 차원 `x`에서 분할됩니다.\n",
        "\n",
        "\n",
        "<img src=\"https://www.tensorflow.org/images/dtensor/dtensor_layout_hybrid.png\" class=\"no-filter\" alt=\"메시 차원 y에서 복제된 텐서, 메시 차원 'x'에서 분할된 첫 번째 축\">      \n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IkWe6mVl7uRb"
      },
      "outputs": [],
      "source": [
        "layout = dtensor.Layout([\"x\", dtensor.UNSHARDED], mesh_2d)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TTalu6M-ISYb"
      },
      "source": [
        "### 단일 클라이언트 및 멀티 클라이언트 애플리케이션\n",
        "\n",
        "DTensor는 단일 클라이언트 및 멀티 클라이언트 애플리케이션을 모두 지원합니다. Colab Python 커널은 단일 Python 프로세스가 있는 단일 클라이언트 DTensor 애플리케이션의 예입니다.\n",
        "\n",
        "멀티 클라이언트 DTensor 애플리케이션에서 멀티 Python 프로세스는 하나의 애플리케이션처럼 공동으로 작업을 수행합니다. 멀티 클라이언트 DTensor 애플리케이션에서 `Mesh`의 데카르트 그리드는 현재 클라이언트에 로컬로 연결되어 있는지 아니면 다른 클라이언트에 원격으로 연결되어 있는지에 관계 없이 여러 장치에 걸쳐 있을 수 있습니다. `Mesh`에서 사용하는 모든 기기 세트는 *전역 기기 목록*이라고 부릅니다.\n",
        "\n",
        "멀티 클라이언트 DTensor 애플리케이션에서 `Mesh`를 생성하는 것은 *전역 기기 목록*이 참여하는 모든 클라이언트에 대해 동일하고 `Mesh`의 생성이 글로벌 장벽 역할을 하는 집합적 연산입니다.\n",
        "\n",
        "`Mesh`를 생성하는 동안 각 클라이언트는 예상 *글로벌 장치 목록*과 함께 *로컬 기기 목록*을 제공합니다. D텐서는 두 목록이 일치하는지 확인합니다. 멀티 클라이언트 메시 생성 및 *글로벌 장치 목록*에 대한 자세한 내용은 `dtensor.create_mesh` 및 `dtensor.create_distributed_mesh`에 대한 API 문서를 참조합니다.\n",
        "\n",
        "단일 클라이언트는 1개의 클라이언트가 있는 멀티 클라이언트의 특수한 경우로 생각할 수 있습니다. 단일 클라이언트 애플리케이션에서 *전역 기기 목록*은 *로컬 기기 목록*과 동일합니다.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P_F7DWkXkB4w"
      },
      "source": [
        "## 분할 텐서로서의 DTensor\n",
        "\n",
        "이제 `DTensor`를 사용하여 코딩을 시작하겠습니다. 도우미 함수인 `dtensor_from_array`는 `tf.Tensor`처럼 보이는 항목으로부터 DTensor를 생성하는 방법을 보여줍니다. 이 함수는 2개의 단계를 수행합니다.\n",
        "\n",
        "- 메시의 모든 기기에 텐서를 복제합니다.\n",
        "- 인수에서 요청한 레이아웃에 따라 사본을 분할합니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "s6aws-b8dN9L"
      },
      "outputs": [],
      "source": [
        "def dtensor_from_array(arr, layout, shape=None, dtype=None):\n",
        "  \"\"\"Convert a DTensor from something that looks like an array or Tensor.\n",
        "\n",
        "  This function is convenient for quick doodling DTensors from a known,\n",
        "  unsharded data object in a single-client environment. This is not the\n",
        "  most efficient way of creating a DTensor, but it will do for this\n",
        "  tutorial.\n",
        "  \"\"\"\n",
        "  if shape is not None or dtype is not None:\n",
        "    arr = tf.constant(arr, shape=shape, dtype=dtype)\n",
        "\n",
        "  # replicate the input to the mesh\n",
        "  a = dtensor.copy_to_mesh(arr,\n",
        "          layout=dtensor.Layout.replicated(layout.mesh, rank=layout.rank))\n",
        "  # shard the copy to the desirable layout\n",
        "  return dtensor.relayout(a, layout=layout)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r3o6IysrlGMu"
      },
      "source": [
        "### DTensor의 해부학\n",
        "\n",
        "DTensor는 `tf.Tensor` 객체이지만 분할 동작을 정의하는 `Layout` 주석으로 확대됩니다. DTensor는 다음으로 구성되어 있습니다.\n",
        "\n",
        "- 텐서의 전역 형상 및 dtype을 포함하는 전역 텐서 메타데이터입니다.\n",
        "- `Tensor`가 속한 `Mesh`를 정의하고 `Tensor`가 `Mesh`에 분할되는 방식인 `Layout`입니다.\n",
        "- **구성 요소 텐서** 목록으로 `Mesh`의 로컬 기기당 하나의 항목만 있습니다.\n",
        "\n",
        "`dtensor from_array`를 사용하여 첫 번째 DTensor인 `my_first_dtensor`를 생성하고 해당 콘텐츠를 검사할 수 있습니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mQu_nScGUvYH"
      },
      "outputs": [],
      "source": [
        "mesh = dtensor.create_mesh([(\"x\", 6)], devices=DEVICES)\n",
        "layout = dtensor.Layout([dtensor.UNSHARDED], mesh)\n",
        "\n",
        "my_first_dtensor = dtensor_from_array([0, 1], layout)\n",
        "\n",
        "# Examine the dtensor content\n",
        "print(my_first_dtensor)\n",
        "print(\"global shape:\", my_first_dtensor.shape)\n",
        "print(\"dtype:\", my_first_dtensor.dtype)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r8LQy1nqmvFy"
      },
      "source": [
        "#### 레이아웃과 `fetch_layout`\n",
        "\n",
        "DTensor의 레이아웃은 `tf.Tensor`의 일반 속성이 아닙니다. 대신 DTensor는 DTensor의 레이아웃에 액세스할 수 있는 `dtensor.fetch_layout` 함수를 제공합니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dCSFyaAjmzGu"
      },
      "outputs": [],
      "source": [
        "print(dtensor.fetch_layout(my_first_dtensor))\n",
        "assert layout == dtensor.fetch_layout(my_first_dtensor)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ed7i3l2lmatm"
      },
      "source": [
        "#### 구성 요소 텐서, `pack`과 `unpack`\n",
        "\n",
        "DTensor는 **구성 요소 텐서** 목록으로 구성되어 있습니다. `Mesh`에 있는 기기의 구성 요소 텐서는 이 기기에 저장된 글로벌 DTensor의 일부를 나타내는 `Tensor` 객체입니다.\n",
        "\n",
        "DTensor는 `dtensor.unpack`을 통해 구성 요소 텐서로 압축을 해제할 수 있습니다. `dtensor.unpack`을 사용하여 DTensor의 구성 요소를 검사하고 구성 요소가 `Mesh`의 모든 장치에 있는지 확인할 수 있습니다.\n",
        "\n",
        "글로벌 보기에서 구성 요소 텐서의 위치는 서로 중첩될 수 있습니다. 예를 들어 완전히 복제된 레이아웃의 경우 모든 구성 요소는 글로벌 텐서와 동일한 복제본입니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BGbjqVAOnXMk"
      },
      "outputs": [],
      "source": [
        "for component_tensor in dtensor.unpack(my_first_dtensor):\n",
        "  print(\"Device:\", component_tensor.device, \",\", component_tensor)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-tqIQM52k788"
      },
      "source": [
        "표시된 대로 `my_first_dtensor`는 6개의 장치 모두에 복제된 `[0, 1]`의 텐서입니다."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6By3k-CGn3yv"
      },
      "source": [
        "`dtensor.unpack`의 역연산은 `dtensor.pack`입니다. 구성 요소 텐서는 DTensor로 다시 압축할 수 있습니다.\n",
        "\n",
        "구성 요소는 반환된 DTensor의 순위 및 dtype과 동일한 순위 및 dtype을 가져야 합니다. 그러나 `dtensor.unpack`의 입력인 구성 요소 텐서의 기기 배치에 대한 엄격한 요구사항은 없습니다. 함수는 구성 요소 텐서를 해당 기기에 자동으로 복사합니다.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9lT-6qQwxOgf"
      },
      "outputs": [],
      "source": [
        "packed_dtensor = dtensor.pack(\n",
        "    [[0, 1], [0, 1], [0, 1],\n",
        "     [0, 1], [0, 1], [0, 1]],\n",
        "     layout=layout\n",
        ")\n",
        "print(packed_dtensor)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zvS3autrpK2U"
      },
      "source": [
        "### 메시에 DTensor 분할하기\n",
        "\n",
        "지금까지 dim-1 `Mesh`에 완전히 복제된 1순위 DTensor인 `my_first_dtensor`로 작업을 진행했습니다.\n",
        "\n",
        "다음에서는 dim-2 `Mesh`에 분할된DTensor를 생성하고 검사합니다. 다음 예제에서는 6개의 CPU 장치에서 3x2 `Mesh`를 사용하여 이 작업을 수행합니다. 여기서 메시 차원 `'x'`의 크기는 3개 장치이고 메시 차원`'y'`의 크기는 2개 장치입니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KWb9Ae0VJ-Rc"
      },
      "outputs": [],
      "source": [
        "mesh = dtensor.create_mesh([(\"x\", 3), (\"y\", 2)], devices=DEVICES)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ndSeQSFWKQk9"
      },
      "source": [
        "#### Dim-2 메시에 완전히 분할된 2순위 텐서\n",
        "\n",
        "3x2 2순위 DTensor를 생성하여 `'x'` 메시 차원을 따라 첫 번째 축을 분할하고 `'y'` 메시 차원을 따라 두 번째 축을 분할합니다.\n",
        "\n",
        "- 텐서 형상은 분할된 모든 축을 따라 메시 차원과 같기 때문에 각 장치는 DTensor의 단일 요소를 수신합니다.\n",
        "- 구성 요소 텐서의 순위는 항상 전역 형상의 순위와 동일합니다. DTensor는 구성 요소 텐서와 글로벌 DTensor 사이의 관계를 찾을 수 있도록 정보를 보존하는 간단한 방식으로 이 규칙을 적용합니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ax_ZHouJp1MX"
      },
      "outputs": [],
      "source": [
        "fully_sharded_dtensor = dtensor_from_array(\n",
        "    tf.reshape(tf.range(6), (3, 2)),\n",
        "    layout=dtensor.Layout([\"x\", \"y\"], mesh))\n",
        "\n",
        "for raw_component in dtensor.unpack(fully_sharded_dtensor):\n",
        "  print(\"Device:\", raw_component.device, \",\", raw_component)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zhsLC-NgrC2p"
      },
      "source": [
        "#### Dim-2 메시에 완전히 복제된 2순위 텐서\n",
        "\n",
        "비교를 위해 동일한 dim-2 메시에 완전히 복제된 3x2 2순위 DTensor를 생성합니다.\n",
        "\n",
        "- DTensor가 완전히 복제되기 때문에 각 기기는 3x2 DTensor의 전체 복제본을 수신합니다.\n",
        "- 구성 요소 텐서의 순위는 전역 형상의 순위와 동일합니다. 이 내용은 중요하지 않습니다. 이 경우 구성 요소 텐서의 형상이 전역 형상과 동일하기 때문입니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xmyC6H6Ec90P"
      },
      "outputs": [],
      "source": [
        "fully_replicated_dtensor = dtensor_from_array(\n",
        "    tf.reshape(tf.range(6), (3, 2)),\n",
        "    layout=dtensor.Layout([dtensor.UNSHARDED, dtensor.UNSHARDED], mesh))\n",
        "# Or, layout=tensor.Layout.fully_replicated(mesh, rank=2)\n",
        "\n",
        "for component_tensor in dtensor.unpack(fully_replicated_dtensor):\n",
        "  print(\"Device:\", component_tensor.device, \",\", component_tensor)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KWoyv_oHMzk1"
      },
      "source": [
        "#### Dim-2 메시의 하이브리드 2순위 텐서\n",
        "\n",
        "완전한 분할과 완전한 복제 사이에 있는 무언가는 어떤가요?\n",
        "\n",
        "DTensor를 사용하면 `Layout`이 하이브리드가 될 수 있습니다. 일부 축을 따라 분할되지만 다른 축을 따라 복제됩니다.\n",
        "\n",
        "예를 들어 동일한 3x2 2순위 DTensor를 다음과 같은 방식으로 분할할 수 있습니다.\n",
        "\n",
        "- `'x'` 메시 차원을 따라 분할된 첫 번째 축입니다.\n",
        "- `'y'` 메시 차원을 따라 복제된 두 번째 축입니다.\n",
        "\n",
        "이 분할 구성표를 작성하려면 두 번째 축의 분할 사양을 `'y'`에서 `dtensor.UNSHARDED`로 교체하고 두 번째 축을 따라 복제할 의도를 나타내면 됩니다. 레이아웃 객체는 `Layout(['x', dtensor.UNSHARDED], mesh)`처럼 보일 것입니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DygnbkQ1Lu42"
      },
      "outputs": [],
      "source": [
        "hybrid_sharded_dtensor = dtensor_from_array(\n",
        "    tf.reshape(tf.range(6), (3, 2)),\n",
        "    layout=dtensor.Layout(['x', dtensor.UNSHARDED], mesh))\n",
        "\n",
        "for component_tensor in dtensor.unpack(hybrid_sharded_dtensor):\n",
        "  print(\"Device:\", component_tensor.device, \",\", component_tensor)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "T7FtZ9kQRZgE"
      },
      "source": [
        "생성된 D텐서의 구성 요소 텐서를 검사하고 구성표에 따라 실제로 분할되었는지 확인할 수 있습니다. 표로 상황을 설명하는 것이 도움이 될 수 있습니다.\n",
        "\n",
        "<img src=\"https://www.tensorflow.org/images/dtensor/dtensor_hybrid_mesh.png\" alt=\"A 3x2 hybrid mesh with 6 CPUs\" class=\"\">\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "auAkA38XjL-q"
      },
      "source": [
        "#### Tensor.numpy() 및 분할된 DTensor\n",
        "\n",
        "분할된 DTensor에서 `.numpy()` 메서드를 호출하면 오류가 발생한다는 점에 유의해야 합니다. 이러한 오류는 여러 컴퓨팅 장치에서 여러 컴퓨팅 장치로부터 호스트 CPU 기기로 반환된 numpy 배열을 지원하며 의도하지 않은 데이터 수집이 수행되는 것을 방지하기 위함입니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hNdwmnL0jAXS"
      },
      "outputs": [],
      "source": [
        "print(fully_replicated_dtensor.numpy())\n",
        "\n",
        "try:\n",
        "  fully_sharded_dtensor.numpy()\n",
        "except tf.errors.UnimplementedError:\n",
        "  print(\"got an error as expected for fully_sharded_dtensor\")\n",
        "\n",
        "try:\n",
        "  hybrid_sharded_dtensor.numpy()\n",
        "except tf.errors.UnimplementedError:\n",
        "  print(\"got an error as expected for hybrid_sharded_dtensor\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8WcMkiagPF_6"
      },
      "source": [
        "## DTensor의 TensorFlow API\n",
        "\n",
        "DTensor는 프로그램에서 텐서를 언제든지 대체하기 위해 노력합니다. TensorFlow Python API는 연산 라이브러리 함수인 `tf.function`, `tf.GradientTape`와 같은 `tf.Tensor`를 사용하며 또한 DTensor와도 작동합니다.\n",
        "\n",
        "이를 수행하기 위해 DTensor는 각 [TensorFlow 그래프](https://www.tensorflow.org/guide/intro_to_graphs)에 대해 *SPMD 확장*이라 부르는 절차에서 동일한 [SPMD](https://en.wikipedia.org/wiki/SPMD) 그래프를 생성하고 실행합니다. DTensor SPMD 확장의 몇 가지 중요한 단계는 다음과 같습니다.\n",
        "\n",
        "- TensorFlow 그래프에서 DTensor의 분할 `Layout` 전파\n",
        "- 글로벌 DTensor의 TensorFlow 연산을 구성 요소 텐서의 동일한 TensorFlow 연산으로 다시 작성하고 필요한 경우 집합 및 통신 연산을 삽입\n",
        "- 백엔드 중립 TensorFlow 연산을 백엔드 특정 TensorFlow 연산으로 낮춤\n",
        "\n",
        "최종 결과는 **DTensor가 텐서의 드롭인 교체품**이라는 것입니다.\n",
        "\n",
        "참고: DTensor는 여전히 실험적인 API이므로 DTensor 프로그래밍 모델의 범위와 한계를 탐구하고 확장하게 됩니다.\n",
        "\n",
        "두 가지 방법을 사용하여 DTensor 실행을 트리거할 수 있습니다.\n",
        "\n",
        "- `tf.matmul(a, b)`와 같은 Python 함수의 피연산자로서의 DTensor는 `a`, `b` 또는 둘 다 DTensor인 경우 DTensor를 통해 실행됩니다.\n",
        "- `dtensor.call_with_layout(tf.ones, layout, shape=(3, 2))`와 같이 Python 함수의 결과가 DTensor가 되도록 요청하면 tf.ones의 출력이 `layout`에 따라 분할되도록 요청했기 때문에 DTensor를 통해 실행됩니다."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "urKzmqAoPssT"
      },
      "source": [
        "### 피연산자로서의 DTensor\n",
        "\n",
        "많은 TensorFlow API 함수는 `tf.Tensor`를 피연산자로 사용하고 그 결과로 `tf.Tensor`를 반환합니다. 이러한 함수의 경우 DTensor를 피연산자로 전달하여 DTensor를 통해 함수를 실행하려는 의도를 표현할 수 있습니다. 이 섹션에서는 `tf.matmul(a, b)`을 예제로 사용합니다."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7LO8ZT7iWVga"
      },
      "source": [
        "#### 완전히 복제된 입력과 출력\n",
        "\n",
        "이 경우 DTensor가 완전히 복제됩니다. `Mesh`의 각 장치에서,\n",
        "\n",
        "- 피연산자 `a`의 구성 요소 텐서는 `[[1, 2, 3], [4, 5, 6]]`(2x3)입니다.\n",
        "- 피연산자 `b`의 구성 요소 텐서는 `[[6, 5], [4, 3], [2, 1]]`(3x2)입니다.\n",
        "- 계산은 `(2x3, 3x2) -> 2x2`의 단일 `MatMul`로 구성됩니다.\n",
        "- 결과 `c`의 구성 요소 텐서는 `[[20, 14], [56,41]]` (2x2)입니다.\n",
        "\n",
        "전체 부동 소수점 mul 연산의 수는 `6 device * 4 result * 3 mul = 72`입니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TiZf2J9JNd2D"
      },
      "outputs": [],
      "source": [
        "mesh = dtensor.create_mesh([(\"x\", 6)], devices=DEVICES)\n",
        "layout = dtensor.Layout([dtensor.UNSHARDED, dtensor.UNSHARDED], mesh)\n",
        "a = dtensor_from_array([[1, 2, 3], [4, 5, 6]], layout=layout)\n",
        "b = dtensor_from_array([[6, 5], [4, 3], [2, 1]], layout=layout)\n",
        "\n",
        "c = tf.matmul(a, b) # runs 6 identical matmuls in parallel on 6 devices\n",
        "\n",
        "# `c` is a DTensor replicated on all devices (same as `a` and `b`)\n",
        "print('Sharding spec:', dtensor.fetch_layout(c).sharding_specs)\n",
        "print(\"components:\")\n",
        "for component_tensor in dtensor.unpack(c):\n",
        "  print(component_tensor.device, component_tensor.numpy())\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QXtR9qgKWgWV"
      },
      "source": [
        "#### 축약된 축을 따라 피연산자 분할하기\n",
        "\n",
        "피연산자 `a` 및 `b`를 분할하여 기기당 계산량을 줄일 수 있습니다. `tf.matmul`에서 많이 사용하는 분할 체계는 축약 축을 따라 피연산자를 분할하는 것입니다. 즉, 두 번째 축을 따라 `a`를, 그리고 첫 번째 축을 따라 `b`를 분할하는 것을 의미합니다.\n",
        "\n",
        "이 체계에 따라 분할된 글로벌 행렬 곱셈은 동시에 실행되는 로컬 matmul에 의해 효율적으로 수행될 수 있으며 로컬 결과를 집계할 수 있도록 집단 축소가 이어서 수행됩니다. 이는 분산 행렬 점 곱셈을 구현하는 [정식 방법](https://github.com/open-mpi/ompi/blob/ee87ec391f48512d3718fc7c8b13596403a09056/docs/man-openmpi/man3/MPI_Reduce.3.rst?plain=1#L265)이기도 합니다.\n",
        "\n",
        "부동 소수점 mul 연산의 총 개수는 `6 devices * 4 result * 1 = 24`이며, 위의 완전히 복제된 경우인 (72)에 비해 3배 감소합니다. 3배 감소하는 원인은 `3`개 기기의 크기를 사용하는 `x` 메시 차원을 따라 분할이 수행되기 때문입니다.\n",
        "\n",
        "순차적으로 실행되는 연산의 수를 줄이는 것이 동기식 모델 병렬 처리가 훈련을 가속화하는 주요 메커니즘입니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EyVAUvMePbms"
      },
      "outputs": [],
      "source": [
        "mesh = dtensor.create_mesh([(\"x\", 3), (\"y\", 2)], devices=DEVICES)\n",
        "a_layout = dtensor.Layout([dtensor.UNSHARDED, 'x'], mesh)\n",
        "a = dtensor_from_array([[1, 2, 3], [4, 5, 6]], layout=a_layout)\n",
        "b_layout = dtensor.Layout(['x', dtensor.UNSHARDED], mesh)\n",
        "b = dtensor_from_array([[6, 5], [4, 3], [2, 1]], layout=b_layout)\n",
        "\n",
        "c = tf.matmul(a, b)\n",
        "# `c` is a DTensor replicated on all devices (same as `a` and `b`)\n",
        "print('Sharding spec:', dtensor.fetch_layout(c).sharding_specs)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IhD8yYgJiCEh"
      },
      "source": [
        "#### 추가 분할\n",
        "\n",
        "입력에 대해 추가 분할을 수행할 수 있으며 입력값은 결과에 적절하게 이전됩니다. 예를 들어 첫 번째 축을 따라 피연산자 `a`의 추가 분할을 `'y'` 메시 차원에 적용할 수 있습니다. 추가 분할은 결과 `c`의 첫 번째 축으로 이전됩니다.\n",
        "\n",
        "부동 소수점 mul 연산의 총 개수는 `6 devices * 2 result * 1 = 12`이며, 위의 경우 (24)에 비해 추가적으로 2배 감소합니다. 2배 감소된 원인은 `2`개 기기의 크기를 사용하는 `y` 메시 차원을 따라 분할이 수행되기 때문입니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0PYqe0neiOpR"
      },
      "outputs": [],
      "source": [
        "mesh = dtensor.create_mesh([(\"x\", 3), (\"y\", 2)], devices=DEVICES)\n",
        "\n",
        "a_layout = dtensor.Layout(['y', 'x'], mesh)\n",
        "a = dtensor_from_array([[1, 2, 3], [4, 5, 6]], layout=a_layout)\n",
        "b_layout = dtensor.Layout(['x', dtensor.UNSHARDED], mesh)\n",
        "b = dtensor_from_array([[6, 5], [4, 3], [2, 1]], layout=b_layout)\n",
        "\n",
        "c = tf.matmul(a, b)\n",
        "# The sharding of `a` on the first axis is carried to `c'\n",
        "print('Sharding spec:', dtensor.fetch_layout(c).sharding_specs)\n",
        "print(\"components:\")\n",
        "for component_tensor in dtensor.unpack(c):\n",
        "  print(component_tensor.device, component_tensor.numpy())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c-1NazCVmLWZ"
      },
      "source": [
        "### 출력으로서의 DTensor\n",
        "\n",
        "피연산자를 사용하지 않지만 분할수 있는 텐서 결과를 반환하는 Python 함수는 어떨까요? 이러한 함수의 예는 다음과 같습니다.\n",
        "\n",
        "- `tf.ones`, `tf.zeros`, `tf.random.stateless_normal`,\n",
        "\n",
        "DTensor는 이러한 Python 함수에 DTensor로 Python 함수를 실행하는 `dtensor.call_with_layout`을 즉시 제공하며, 반환된 텐서가 요청한 `Layout`이 있는 DTensor인지 확인합니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "J0jo_8NPtJiO"
      },
      "outputs": [],
      "source": [
        "help(dtensor.call_with_layout)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "V-YdLvfytM7g"
      },
      "source": [
        "즉시 실행되는 Python 함수는 일반적으로 중요하지 않은 단일 TensorFlow 연산만 포함합니다.\n",
        "\n",
        "`dtensor.call_with_layout`가 있는 멀티 TensorFlow 연산을 방출하는 Python 함수를 사용하려면 함수를 `tf.function`으로 변환해야 합니다. `tf.function` 호출은 단일 TensorFlow 연산입니다. `tf.function`이 호출되는 경우 DTensor는 중간 텐서가 구체화되기 전에 `tf.function`의 계산 그래프를 분석할 때 레이아웃 전파를 수행할 수 있습니다."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DLrksgFjqRLS"
      },
      "source": [
        "#### 단일 TensorFlow 연산을 방출하는 API\n",
        "\n",
        "함수가 단일 TensorFlow 연산을 방출하는 경우 `dtensor.call_with_layout`을 함수에 직접 적용할 수 있습니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "G1CuKYSFtFeM"
      },
      "outputs": [],
      "source": [
        "help(tf.ones)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2m_EAwy-ozOh"
      },
      "outputs": [],
      "source": [
        "mesh = dtensor.create_mesh([(\"x\", 3), (\"y\", 2)], devices=DEVICES)\n",
        "ones = dtensor.call_with_layout(tf.ones, dtensor.Layout(['x', 'y'], mesh), shape=(6, 4))\n",
        "print(ones)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bx-7Xo8Cpb8S"
      },
      "source": [
        "#### 멀티 TensorFlow 연산을 방출하는 API\n",
        "\n",
        "API가 멀티 TensorFlow 연산을 방출하는 경우 `tf.function`을 통해 함수를 단일 연산으로 변환합니다. `tf.random.stateleess_normal`이 그 예입니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "H8BQSTRFtCih"
      },
      "outputs": [],
      "source": [
        "help(tf.random.stateless_normal)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TvP81eYopSPm"
      },
      "outputs": [],
      "source": [
        "ones = dtensor.call_with_layout(\n",
        "    tf.function(tf.random.stateless_normal),\n",
        "    dtensor.Layout(['x', 'y'], mesh),\n",
        "    shape=(6, 4),\n",
        "    seed=(1, 1))\n",
        "print(ones)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qKoojp9ZyWzW"
      },
      "source": [
        "단일 TensorFlow 연산을 방출하는 Python 함수를 `tf.function`으로 래핑할 수 있습니다. 유일한 주의 사항은 Python 함수에서 `tf.function`을 생성할 때 관련 비용과 복잡성을 신경써야 한다는 것입니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LbAtKrSkpOaq"
      },
      "outputs": [],
      "source": [
        "ones = dtensor.call_with_layout(\n",
        "    tf.function(tf.ones),\n",
        "    dtensor.Layout(['x', 'y'], mesh),\n",
        "    shape=(6, 4))\n",
        "print(ones)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D-m1816JP3CE"
      },
      "source": [
        "### `tf.Variable`에서 `tensor.Variable`로\n",
        "\n",
        "Tensorflow에서 `tf.Variable`은 변경 가능한 `Tensor` 값의 홀더입니다. DTensor를 사용할 경우 `dtensor.DVariable`에 의해 해당 변수 의미 체계가 제공됩니다.\n",
        "\n",
        "DTensor 변수에 새로운 유형의 `Variable`이 도입된 이유는 D변수(DVariables)에 레이아웃이 초기 값에서 변경할 수 없다는 추가 요구 사항이 있기 때문입니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "awRPuR26P0Sc"
      },
      "outputs": [],
      "source": [
        "mesh = dtensor.create_mesh([(\"x\", 6)], devices=DEVICES)\n",
        "layout = dtensor.Layout([dtensor.UNSHARDED, dtensor.UNSHARDED], mesh)\n",
        "\n",
        "v = dtensor.DVariable(\n",
        "    initial_value=dtensor.call_with_layout(\n",
        "        tf.function(tf.random.stateless_normal),\n",
        "        layout=layout,\n",
        "        shape=tf.TensorShape([64, 32]),\n",
        "        seed=[1, 1],\n",
        "        dtype=tf.float32))\n",
        "\n",
        "print(v.handle)\n",
        "assert layout == dtensor.fetch_layout(v)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Pb9jn473prC_"
      },
      "source": [
        "`layout` 일치에 대한 요구사항을 제외하고 `DVariable`은 `tf.Variable`과 동일하게 작동합니다. 예를 들어, DTensor에 D변수(DVariable)를 추가할 수 있습니다.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "adxFw9wJpqQQ"
      },
      "outputs": [],
      "source": [
        "a = dtensor.call_with_layout(tf.ones, layout=layout, shape=(64, 32))\n",
        "b = v + a # add DVariable and DTensor\n",
        "print(b)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QxBdNHWSu-kV"
      },
      "source": [
        "또한 D변수에 DTensor를 할당할 수도 있습니다.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oYwfiyw5P94U"
      },
      "outputs": [],
      "source": [
        "v.assign(a) # assign a DTensor to a DVariable\n",
        "print(a)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4fvSk_VUvGnj"
      },
      "source": [
        "호환되지 않는 레이아웃이 있는 DTensor를 할당하여 `DVariable`의 레이아웃을 변경하려고 하면 오류가 발생합니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3pckUugYP_r-"
      },
      "outputs": [],
      "source": [
        "# variable's layout is immutable.\n",
        "another_mesh = dtensor.create_mesh([(\"x\", 3), (\"y\", 2)], devices=DEVICES)\n",
        "b = dtensor.call_with_layout(tf.ones,\n",
        "                     layout=dtensor.Layout([dtensor.UNSHARDED, dtensor.UNSHARDED], another_mesh),\n",
        "                     shape=(64, 32))\n",
        "try:\n",
        "  v.assign(b)\n",
        "except:\n",
        "  print(\"exception raised\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3LadIcwRvR6f"
      },
      "source": [
        "## 다음 단계\n",
        "\n",
        "이 Colab에서는 분산 컴퓨팅에 대한 TensorFlow의 확장 개념인 DTensor에 대해 배웠습니다. 튜토리얼에서 이 개념을 사용하려면 [DTensor를 사용한 분산 훈련](https://www.tensorflow.org/tutorials/distribute/dtensor_ml_tutorial)을 참고하세요."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "dtensor_overview.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}