{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "16749a86",
   "metadata": {},
   "source": [
    "# SPMD JAX on IPUs: `pjit` quickstart"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f433e95c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found existing installation: jax 0.3.16+ipu\n",
      "Uninstalling jax-0.3.16+ipu:\n",
      "  Successfully uninstalled jax-0.3.16+ipu\n",
      "Found existing installation: jaxlib 0.3.15+ipu.sdk310\n",
      "Uninstalling jaxlib-0.3.15+ipu.sdk310:\n",
      "  Successfully uninstalled jaxlib-0.3.15+ipu.sdk310\n",
      "Looking in links: https://graphcore-research.github.io/jax-experimental/wheels.html\n",
      "Collecting jax==0.3.16+ipu\n",
      "  Using cached https://github.com/graphcore-research/jax-experimental/releases/download/jax-v0.3.16-ipu-beta2-sdk3/jax-0.3.16%2Bipu-py3-none-any.whl (1.2 MB)\n",
      "Collecting jaxlib==0.3.15+ipu.sdk310\n",
      "  Using cached https://github.com/graphcore-research/jax-experimental/releases/download/jax-v0.3.16-ipu-beta2-sdk3/jaxlib-0.3.15%2Bipu.sdk310-cp38-none-manylinux2014_x86_64.whl (109.4 MB)\n",
      "Requirement already satisfied: absl-py in /nethome/paulb/venvs/3.1.0+1205/3.1.0+1205_poptorch/lib/python3.8/site-packages (from jax==0.3.16+ipu) (1.4.0)\n",
      "Requirement already satisfied: numpy>=1.20 in /nethome/paulb/venvs/3.1.0+1205/3.1.0+1205_poptorch/lib/python3.8/site-packages (from jax==0.3.16+ipu) (1.22.4)\n",
      "Requirement already satisfied: opt-einsum in /nethome/paulb/venvs/3.1.0+1205/3.1.0+1205_poptorch/lib/python3.8/site-packages (from jax==0.3.16+ipu) (3.3.0)\n",
      "Requirement already satisfied: scipy>=1.5 in /nethome/paulb/venvs/3.1.0+1205/3.1.0+1205_poptorch/lib/python3.8/site-packages (from jax==0.3.16+ipu) (1.10.0)\n",
      "Requirement already satisfied: typing-extensions in /nethome/paulb/venvs/3.1.0+1205/3.1.0+1205_poptorch/lib/python3.8/site-packages (from jax==0.3.16+ipu) (4.4.0)\n",
      "Requirement already satisfied: etils[epath] in /nethome/paulb/venvs/3.1.0+1205/3.1.0+1205_poptorch/lib/python3.8/site-packages (from jax==0.3.16+ipu) (1.0.0)\n",
      "Requirement already satisfied: importlib_resources in /nethome/paulb/venvs/3.1.0+1205/3.1.0+1205_poptorch/lib/python3.8/site-packages (from etils[epath]->jax==0.3.16+ipu) (5.10.2)\n",
      "Requirement already satisfied: zipp in /nethome/paulb/venvs/3.1.0+1205/3.1.0+1205_poptorch/lib/python3.8/site-packages (from etils[epath]->jax==0.3.16+ipu) (3.12.0)\n",
      "Installing collected packages: jaxlib, jax\n",
      "Successfully installed jax-0.3.16+ipu jaxlib-0.3.15+ipu.sdk310\n"
     ]
    }
   ],
   "source": [
    "# Install experimental JAX for IPUs (SDK 3.1) from Github releases.\n",
    "import sys\n",
    "!{sys.executable} -m pip uninstall -y jax jaxlib\n",
    "!{sys.executable} -m pip install jax==0.3.16+ipu jaxlib==0.3.15+ipu.sdk310 -f https://graphcore-research.github.io/jax-experimental/wheels.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "126bc13d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax.config import config\n",
    "\n",
    "# Select how many IPUs will be visible.\n",
    "config.FLAGS.jax_ipu_device_count = 4\n",
    "\n",
    "# Simulating `pmap` on CPU devices instead.\n",
    "# os.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=4'\n",
    "# config.FLAGS.jax_platforms = \"cpu,ipu\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "cbf2f003",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "329c05e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "from jax.experimental import maps\n",
    "from jax.experimental import PartitionSpec\n",
    "from jax.experimental.pjit import pjit\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ea5b3f7b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "IPU devices:\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[IpuDevice(id=0, num_tiles=1472, version=ipu2),\n",
       " IpuDevice(id=1, num_tiles=1472, version=ipu2),\n",
       " IpuDevice(id=2, num_tiles=1472, version=ipu2),\n",
       " IpuDevice(id=3, num_tiles=1472, version=ipu2)]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(\"IPU devices:\")\n",
    "ipu_devices = jax.devices(\"ipu\")\n",
    "ipu_devices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d62f0657",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "34df81f1",
   "metadata": {},
   "source": [
    "# Basic `pjit` examples"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b179af8",
   "metadata": {},
   "source": [
    "## Mesh definition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "fb09b096",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "IPU device mesh: Mesh(array([[0, 1],\n",
      "       [2, 3]]), ('x', 'y'))\n"
     ]
    }
   ],
   "source": [
    "mesh_shape = (2, 2)\n",
    "mesh_size = np.prod(mesh_shape)\n",
    "mesh_devices = np.asarray(ipu_devices).reshape(*mesh_shape)\n",
    "# 'x', 'y' axis names are used here for simplicity\n",
    "mesh = maps.Mesh(mesh_devices, ('x', 'y'))\n",
    "print(\"IPU device mesh:\", mesh)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "76cfb42c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INPUT:\n",
      " [[ 0.  1.]\n",
      " [ 2.  3.]\n",
      " [ 4.  5.]\n",
      " [ 6.  7.]\n",
      " [ 8.  9.]\n",
      " [10. 11.]] (6, 2)\n"
     ]
    }
   ],
   "source": [
    "# Input data to shard or replicate on IPU mesh.\n",
    "N = 3\n",
    "input_data = np.arange(N * mesh_size, dtype=np.float32).reshape(-1, 2)\n",
    "print(\"INPUT:\\n\", input_data, input_data.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00e51fb5",
   "metadata": {},
   "source": [
    "## `x` and `y` axes output sharding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "2e2a74e1",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: The compile time engine option debug.branchRecordTile is set to \"5887\" when creating the Engine. (At compile-tile it was set to 1471)\n"
     ]
    }
   ],
   "source": [
    "# No partition for inputs: data is replicated.\n",
    "in_axis_resources = None\n",
    "# Output is sharded on x and y axes.\n",
    "out_axis_resources=PartitionSpec('x', 'y')\n",
    "\n",
    "f = pjit(\n",
    "    # Simple unary op to run\n",
    "    lambda x: jax.lax.integer_pow(x, 2),\n",
    "    in_axis_resources=in_axis_resources,\n",
    "    out_axis_resources=out_axis_resources)\n",
    " \n",
    "# Sends data to accelerators based on partition_spec\n",
    "with maps.Mesh(mesh.devices, mesh.axis_names):\n",
    "    output_data = f(input_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c46bc14e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(ShardedDeviceArray([[  0.,   1.],\n",
       "                     [  4.,   9.],\n",
       "                     [ 16.,  25.],\n",
       "                     [ 36.,  49.],\n",
       "                     [ 64.,  81.],\n",
       "                     [100., 121.]], dtype=float32),\n",
       " [DeviceArray([[ 0.],\n",
       "               [ 4.],\n",
       "               [16.]], dtype=float32),\n",
       "  DeviceArray([[ 1.],\n",
       "               [ 9.],\n",
       "               [25.]], dtype=float32),\n",
       "  DeviceArray([[ 36.],\n",
       "               [ 64.],\n",
       "               [100.]], dtype=float32),\n",
       "  DeviceArray([[ 49.],\n",
       "               [ 81.],\n",
       "               [121.]], dtype=float32)])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Full output, and how it is sharded between devices\n",
    "output_data, output_data.device_buffers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e11f801b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "07d43ae0",
   "metadata": {},
   "source": [
    "## `x` axis output sharding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "89c7369c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: The compile time engine option debug.branchRecordTile is set to \"5887\" when creating the Engine. (At compile-tile it was set to 1471)\n"
     ]
    }
   ],
   "source": [
    "# No partition for inputs: data is replicated.\n",
    "in_axis_resources = None\n",
    "# Output is sharded on x axes.\n",
    "out_axis_resources=PartitionSpec('x', None)\n",
    "\n",
    "f = pjit(\n",
    "    # Simple unary op to run\n",
    "    lambda x: jax.lax.integer_pow(x, 2),\n",
    "    in_axis_resources=in_axis_resources,\n",
    "    out_axis_resources=out_axis_resources)\n",
    " \n",
    "# Sends data to accelerators based on partition_spec\n",
    "with maps.Mesh(mesh.devices, mesh.axis_names):\n",
    "    output_data = f(input_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "06ea345d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(ShardedDeviceArray([[  0.,   1.],\n",
       "                     [  4.,   9.],\n",
       "                     [ 16.,  25.],\n",
       "                     [ 36.,  49.],\n",
       "                     [ 64.,  81.],\n",
       "                     [100., 121.]], dtype=float32),\n",
       " [DeviceArray([[ 0.,  1.],\n",
       "               [ 4.,  9.],\n",
       "               [16., 25.]], dtype=float32),\n",
       "  DeviceArray([[ 0.,  1.],\n",
       "               [ 4.,  9.],\n",
       "               [16., 25.]], dtype=float32),\n",
       "  DeviceArray([[ 36.,  49.],\n",
       "               [ 64.,  81.],\n",
       "               [100., 121.]], dtype=float32),\n",
       "  DeviceArray([[ 36.,  49.],\n",
       "               [ 64.,  81.],\n",
       "               [100., 121.]], dtype=float32)])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Output: sharded along X axis, replicated along Y.\n",
    "output_data, output_data.device_buffers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b6df15e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "77e840e5",
   "metadata": {},
   "source": [
    "## `y` axis output sharding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "0ae5a183",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: The compile time engine option debug.branchRecordTile is set to \"5887\" when creating the Engine. (At compile-tile it was set to 1471)\n"
     ]
    }
   ],
   "source": [
    "# No partition for inputs: data is replicated.\n",
    "in_axis_resources = None\n",
    "# Output is sharded on y axes.\n",
    "out_axis_resources=PartitionSpec('y', None)\n",
    "\n",
    "f = pjit(\n",
    "    # Simple unary op to run\n",
    "    lambda x: jax.lax.integer_pow(x, 2),\n",
    "    in_axis_resources=in_axis_resources,\n",
    "    out_axis_resources=out_axis_resources)\n",
    " \n",
    "# Sends data to accelerators based on partition_spec\n",
    "with maps.Mesh(mesh.devices, mesh.axis_names):\n",
    "    output_data = f(input_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "5bdee85b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(ShardedDeviceArray([[  0.,   1.],\n",
       "                     [  4.,   9.],\n",
       "                     [ 16.,  25.],\n",
       "                     [ 36.,  49.],\n",
       "                     [ 64.,  81.],\n",
       "                     [100., 121.]], dtype=float32),\n",
       " [DeviceArray([[ 0.,  1.],\n",
       "               [ 4.,  9.],\n",
       "               [16., 25.]], dtype=float32),\n",
       "  DeviceArray([[ 36.,  49.],\n",
       "               [ 64.,  81.],\n",
       "               [100., 121.]], dtype=float32),\n",
       "  DeviceArray([[ 0.,  1.],\n",
       "               [ 4.,  9.],\n",
       "               [16., 25.]], dtype=float32),\n",
       "  DeviceArray([[ 36.,  49.],\n",
       "               [ 64.,  81.],\n",
       "               [100., 121.]], dtype=float32)])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Output: sharded along Y axis, replicated along X.\n",
    "output_data, output_data.device_buffers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c32bb2b2",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "e10043ae",
   "metadata": {},
   "source": [
    "## (`x`,`y`) axis output sharding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "fec7a3a9",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: The compile time engine option debug.branchRecordTile is set to \"5887\" when creating the Engine. (At compile-tile it was set to 1471)\n"
     ]
    }
   ],
   "source": [
    "# Requires first axis divisible by mesh size!\n",
    "input_data = np.arange(mesh_size * mesh_size, dtype=np.float32).reshape(-1, 2)\n",
    "\n",
    "# No partition for inputs: data is replicated.\n",
    "in_axis_resources = None\n",
    "# Output is sharded on x+y axes.\n",
    "out_axis_resources=PartitionSpec(('x', 'y'), None)\n",
    "\n",
    "f = pjit(\n",
    "    # Simple unary op to run\n",
    "    lambda x: jax.lax.integer_pow(x, 2),\n",
    "    in_axis_resources=in_axis_resources,\n",
    "    out_axis_resources=out_axis_resources)\n",
    " \n",
    "# Sends data to accelerators based on partition_spec\n",
    "with maps.Mesh(mesh.devices, mesh.axis_names):\n",
    "    output_data = f(input_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "d0cab928",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(ShardedDeviceArray([[  0.,   1.],\n",
       "                     [  4.,   9.],\n",
       "                     [ 16.,  25.],\n",
       "                     [ 36.,  49.],\n",
       "                     [ 64.,  81.],\n",
       "                     [100., 121.],\n",
       "                     [144., 169.],\n",
       "                     [196., 225.]], dtype=float32),\n",
       " [DeviceArray([[0., 1.],\n",
       "               [4., 9.]], dtype=float32),\n",
       "  DeviceArray([[16., 25.],\n",
       "               [36., 49.]], dtype=float32),\n",
       "  DeviceArray([[ 64.,  81.],\n",
       "               [100., 121.]], dtype=float32),\n",
       "  DeviceArray([[144., 169.],\n",
       "               [196., 225.]], dtype=float32)])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output_data, output_data.device_buffers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a3d1151",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "37c4fe9b",
   "metadata": {},
   "source": [
    "# Matmul `pjit` example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "33e9e203",
   "metadata": {},
   "outputs": [],
   "source": [
    "M, N, K = 128, 64, 256\n",
    "M, N, K = 12, 8, 16\n",
    "\n",
    "lhs = np.random.rand(M, N).astype(np.float32)\n",
    "rhs = np.random.rand(N, K).astype(np.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "f252222e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: proper sharding of inputs?\n",
    "in_axis_resources = None\n",
    "# Output is sharded on x and y axes.\n",
    "out_axis_resources=PartitionSpec('x', 'y')\n",
    "\n",
    "def compute_fn(lhs, rhs):\n",
    "    return lhs @ rhs\n",
    "\n",
    "f = pjit(\n",
    "    compute_fn,\n",
    "    in_axis_resources=in_axis_resources,\n",
    "    out_axis_resources=out_axis_resources)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "09b2c480",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: The compile time engine option debug.branchRecordTile is set to \"5887\" when creating the Engine. (At compile-tile it was set to 1471)\n"
     ]
    }
   ],
   "source": [
    "# Sends data to accelerators based on partition_spec\n",
    "with maps.Mesh(mesh.devices, mesh.axis_names):\n",
    "    output = f(lhs, rhs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "8bd94cbc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ShardedDeviceArray([[1.5954874 , 1.568143  , 2.1878989 , 2.2802362 ,\n",
       "                     1.9145454 , 1.326968  , 2.38342   , 1.4342963 ,\n",
       "                     1.8710132 , 2.2760766 , 1.2393537 , 1.806525  ,\n",
       "                     1.6943976 , 1.8497736 , 2.0054781 , 1.9822028 ],\n",
       "                    [2.5935414 , 2.467551  , 3.4812937 , 3.4287057 ,\n",
       "                     2.6669626 , 2.2495036 , 3.5671082 , 2.6938055 ,\n",
       "                     2.4568684 , 3.3633797 , 2.2307868 , 3.063488  ,\n",
       "                     2.4313595 , 2.6677954 , 2.6706545 , 3.0482678 ],\n",
       "                    [1.2307926 , 1.468817  , 1.5824169 , 1.6121603 ,\n",
       "                     1.6673343 , 0.95006704, 2.0006866 , 1.3617406 ,\n",
       "                     1.1706672 , 1.3099413 , 1.0319011 , 1.400886  ,\n",
       "                     1.2033337 , 1.3988451 , 0.9909877 , 0.90187824],\n",
       "                    [1.8417573 , 1.8773335 , 1.8307042 , 2.980454  ,\n",
       "                     2.2116432 , 2.26642   , 2.7127883 , 2.1408713 ,\n",
       "                     2.8655107 , 3.3783598 , 2.4769158 , 2.7626133 ,\n",
       "                     1.7139378 , 2.5256588 , 2.8091488 , 3.1348836 ],\n",
       "                    [2.165068  , 1.904313  , 2.5754702 , 3.2189054 ,\n",
       "                     2.3253567 , 2.1346264 , 3.1098194 , 2.4323661 ,\n",
       "                     2.4388776 , 2.8433118 , 2.0695078 , 2.856908  ,\n",
       "                     2.1046596 , 2.5718737 , 2.4854958 , 2.806843  ],\n",
       "                    [1.3993658 , 1.132266  , 1.199956  , 1.6633179 ,\n",
       "                     1.0842177 , 1.2301935 , 1.6186645 , 1.585735  ,\n",
       "                     1.1661737 , 1.5198923 , 1.1500716 , 1.6669728 ,\n",
       "                     0.7827802 , 1.5109612 , 1.1592007 , 1.2610594 ],\n",
       "                    [2.2975454 , 1.6895839 , 2.9678905 , 2.9333122 ,\n",
       "                     1.8801018 , 1.9443885 , 2.8823152 , 2.5690153 ,\n",
       "                     1.9338725 , 2.914373  , 1.841821  , 2.9253213 ,\n",
       "                     1.9866215 , 2.256054  , 2.406379  , 2.9907126 ],\n",
       "                    [2.0274806 , 1.9761821 , 2.3630977 , 2.542418  ,\n",
       "                     1.6337898 , 2.0575063 , 2.2687018 , 1.9539872 ,\n",
       "                     1.8190477 , 2.6270218 , 1.9416327 , 2.1815038 ,\n",
       "                     1.4401137 , 1.9255936 , 1.8207574 , 2.2410314 ],\n",
       "                    [1.7304212 , 1.7692802 , 2.014131  , 2.548343  ,\n",
       "                     1.9773198 , 1.6694659 , 2.5115626 , 1.6842196 ,\n",
       "                     2.0533838 , 2.4655993 , 1.70733   , 2.0088844 ,\n",
       "                     1.6166315 , 2.0828776 , 2.1160207 , 2.1392603 ],\n",
       "                    [2.5761597 , 1.8596513 , 2.617828  , 3.3088527 ,\n",
       "                     1.7990946 , 2.3879538 , 2.8241284 , 2.5008504 ,\n",
       "                     2.4709802 , 3.4961638 , 2.123071  , 2.9737554 ,\n",
       "                     1.7278748 , 2.7171924 , 2.9316914 , 3.2924275 ],\n",
       "                    [1.5923193 , 1.2084457 , 1.8704239 , 2.02193   ,\n",
       "                     1.0130341 , 1.2846794 , 1.7015601 , 1.1669011 ,\n",
       "                     1.1666058 , 1.8143071 , 1.0433682 , 1.3058586 ,\n",
       "                     1.1288259 , 1.4786973 , 1.6271735 , 1.6708196 ],\n",
       "                    [1.5934689 , 1.7857262 , 1.998116  , 2.4237769 ,\n",
       "                     1.923897  , 1.5875312 , 2.3696773 , 1.3797243 ,\n",
       "                     2.0665195 , 2.6399872 , 1.7343158 , 1.7537261 ,\n",
       "                     1.5853028 , 1.9207658 , 2.253568  , 2.2490618 ]],                   dtype=float32)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9adec6e8",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "273aa924",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}