jax_tests.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "flags = os.environ.get('XLA_FLAGS', '')\n",
    "os.environ['XLA_FLAGS'] = flags + \" --xla_force_host_platform_device_count=8\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([11., 20., 29.], dtype=float32)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "\n",
    "x = jnp.arange(5)\n",
    "w = jnp.array([2., 3., 4.])\n",
    "\n",
    "def convolve(x, w):\n",
    "    output = []\n",
    "    for i in range(1, len(x)-1):\n",
    "        output.append(jnp.dot(x[i-1:i+2], w))\n",
    "    return jnp.array(output)\n",
    "convolve(x, w)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "xs = jnp.stack([x, x])\n",
    "ws = jnp.stack([w, w])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[11., 20., 29.],\n",
       "       [11., 20., 29.]], dtype=float32)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "auto_batch_convolve = jax.vmap(convolve)\n",
    "\n",
    "auto_batch_convolve(xs, ws)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[11., 20., 29.],\n",
       "       [11., 20., 29.]], dtype=float32)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch_convolve_v3 = jax.vmap(convolve, in_axes=[0, None])\n",
    "\n",
    "batch_convolve_v3(xs, w)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "We have 8 fake JAX devices now: [CpuDevice(id=0), CpuDevice(id=1), CpuDevice(id=2), CpuDevice(id=3), CpuDevice(id=4), CpuDevice(id=5), CpuDevice(id=6), CpuDevice(id=7)]\n"
     ]
    }
   ],
   "source": [
    "import jax\n",
    "jax.config.update('jax_platform_name', 'cpu')\n",
    "print(f'We have 8 fake JAX devices now: {jax.devices()}')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{CpuDevice(id=0)}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import jax.numpy as jnp\n",
    "\n",
    "arr = jnp.arange(32.0).reshape(4, 8)\n",
    "arr.devices()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[ 0.,  1.,  2.,  3.,  4.,  5.,  6.,  7.],\n",
       "       [ 8.,  9., 10., 11., 12., 13., 14., 15.],\n",
       "       [16., 17., 18., 19., 20., 21., 22., 23.],\n",
       "       [24., 25., 26., 27., 28., 29., 30., 31.]], dtype=float32)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "arr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SingleDeviceSharding(device=CpuDevice(id=0))"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "arr.sharding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                  </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                  </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                  </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                  </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                  </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                      CPU 0                       </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                  </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                  </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                  </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                  </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                  </span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                  \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                  \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                  \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                  \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                  \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m                      \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mCPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m                       \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                  \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                  \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                  \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                  \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                  \u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "jax.debug.visualize_array_sharding(arr)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NamedSharding(mesh=Mesh('x': 2, 'y': 4), spec=PartitionSpec('x', 'y'))\n"
     ]
    }
   ],
   "source": [
    "# Pardon the boilerplate; constructing a sharding will become easier in future!\n",
    "from jax.sharding import Mesh\n",
    "from jax.sharding import PartitionSpec\n",
    "from jax.sharding import NamedSharding\n",
    "from jax.experimental import mesh_utils\n",
    "\n",
    "P = jax.sharding.PartitionSpec\n",
    "devices = mesh_utils.create_device_mesh((2,4))\n",
    "mesh = jax.sharding.Mesh(devices, ('x', 'y'))\n",
    "sharding = jax.sharding.NamedSharding(mesh, P('x', 'y'))\n",
    "print(sharding)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 0.  1.  2.  3.  4.  5.  6.  7.]\n",
      " [ 8.  9. 10. 11. 12. 13. 14. 15.]\n",
      " [16. 17. 18. 19. 20. 21. 22. 23.]\n",
      " [24. 25. 26. 27. 28. 29. 30. 31.]]\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #d6616b\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8ca252\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">            </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #d6616b\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8ca252\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">            </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">   CPU 0    </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #d6616b\">   CPU 1    </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8ca252\">   CPU 2    </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">   CPU 3    </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #d6616b\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8ca252\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">            </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #d6616b\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8ca252\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">            </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #d6616b\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8ca252\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">            </span>\n",
       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #e7cb94\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #6b6ecf\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #a55194\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8c6d31\">            </span>\n",
       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #e7cb94\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #6b6ecf\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #a55194\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8c6d31\">            </span>\n",
       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #e7cb94\">   CPU 4    </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #6b6ecf\">   CPU 5    </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #a55194\">   CPU 6    </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8c6d31\">   CPU 7    </span>\n",
       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #e7cb94\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #6b6ecf\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #a55194\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8c6d31\">            </span>\n",
       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #e7cb94\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #6b6ecf\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #a55194\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8c6d31\">            </span>\n",
       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #e7cb94\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #6b6ecf\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #a55194\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8c6d31\">            </span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m            \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m            \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m            \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m            \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m            \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m            \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m   \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mCPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m    \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m   \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107mCPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m    \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m   \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82mCPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m    \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m   \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mCPU 3\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m    \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m            \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m            \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m            \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m            \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m            \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m            \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m            \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m            \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m            \u001b[0m\n",
       "\u001b[38;2;0;0;0;48;2;231;203;148m            \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m            \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m            \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m            \u001b[0m\n",
       "\u001b[38;2;0;0;0;48;2;231;203;148m            \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m            \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m            \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m            \u001b[0m\n",
       "\u001b[38;2;0;0;0;48;2;231;203;148m   \u001b[0m\u001b[38;2;0;0;0;48;2;231;203;148mCPU 4\u001b[0m\u001b[38;2;0;0;0;48;2;231;203;148m    \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m   \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207mCPU 5\u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m    \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m   \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148mCPU 6\u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m    \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m   \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49mCPU 7\u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m    \u001b[0m\n",
       "\u001b[38;2;0;0;0;48;2;231;203;148m            \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m            \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m            \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m            \u001b[0m\n",
       "\u001b[38;2;0;0;0;48;2;231;203;148m            \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m            \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m            \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m            \u001b[0m\n",
       "\u001b[38;2;0;0;0;48;2;231;203;148m            \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m            \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m            \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m            \u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "arr_sharded = jax.device_put(arr, sharding)\n",
    "\n",
    "print(arr_sharded)\n",
    "jax.debug.visualize_array_sharding(arr_sharded)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "NamedSharding(mesh=Mesh('x': 2, 'y': 4), spec=PartitionSpec('x', 'y'))"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@jax.jit\n",
    "def f_elementwise(x):\n",
    "    return 2 * jnp.sin(x) + 1\n",
    "\n",
    "result = f_elementwise(arr_sharded)\n",
    "\n",
    "result.sharding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[ 1.        ,  2.682942  ,  2.818595  ,  1.28224   , -0.513605  ,\n",
       "        -0.9178486 ,  0.44116902,  2.3139732 ],\n",
       "       [ 2.9787164 ,  1.824237  , -0.08804226, -0.99998045, -0.07314587,\n",
       "         1.840334  ,  2.9812148 ,  2.3005757 ],\n",
       "       [ 0.42419338, -0.92279494, -0.50197446,  1.2997544 ,  2.8258905 ,\n",
       "         2.6733112 ,  0.98229736, -0.69244087],\n",
       "       [-0.81115675,  0.7352965 ,  2.525117  ,  2.912752  ,  1.5418116 ,\n",
       "        -0.32726777, -0.97606325,  0.19192469]], dtype=float32)"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\"> CPU 0,4 </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\"> CPU 1,5 </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\"> CPU 2,6 </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\"> CPU 3,7 </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mCPU 0,4\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mCPU 1,5\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mCPU 2,6\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mCPU 3,7\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[48. 52. 56. 60. 64. 68. 72. 76.]\n"
     ]
    }
   ],
   "source": [
    "@jax.jit\n",
    "def f_contract(x):\n",
    "  return x.sum(axis=0)\n",
    "\n",
    "result = f_contract(arr_sharded)\n",
    "jax.debug.visualize_array_sharding(result)\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">  CPU 0  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #d6616b\">  CPU 1  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8ca252\">  CPU 2  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">  CPU 3  </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #e7cb94\">  CPU 4  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #6b6ecf\">  CPU 5  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #a55194\">  CPU 6  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8c6d31\">  CPU 7  </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #d6616b\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8ca252\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #e7cb94\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #6b6ecf\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #a55194\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8c6d31\">         </span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mCPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m  \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107mCPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m  \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m  \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82mCPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mCPU 3\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\u001b[38;2;0;0;0;48;2;231;203;148m  \u001b[0m\u001b[38;2;0;0;0;48;2;231;203;148mCPU 4\u001b[0m\u001b[38;2;0;0;0;48;2;231;203;148m  \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m  \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207mCPU 5\u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m  \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m  \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148mCPU 6\u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m  \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m  \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49mCPU 7\u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m  \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m         \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;0;0;0;48;2;231;203;148m         \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m         \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m         \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m         \u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "out = jax.pmap(lambda x: x ** 2)\n",
    "\n",
    "result = out(jnp.arange(8))\n",
    "\n",
    "jax.debug.visualize_array_sharding(result)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-06-08 21:44:56.065851: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
      "2024-06-08 21:44:56.066150: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
      "2024-06-08 21:44:56.189252: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
      "2024-06-08 21:44:57.683340: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n"
     ]
    }
   ],
   "source": [
    "import functools\n",
    "from typing import Optional, Callable\n",
    "\n",
    "import numpy as np\n",
    "import jax\n",
    "from jax import lax, random, numpy as jnp\n",
    "\n",
    "import flax\n",
    "from flax import struct, traverse_util, linen as nn\n",
    "from flax.core import freeze, unfreeze\n",
    "from flax.training import train_state, checkpoints\n",
    "\n",
    "import optax # Optax for common losses and optimizers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax.sharding import Mesh, PartitionSpec, NamedSharding\n",
    "from jax.lax import with_sharding_constraint\n",
    "from jax.experimental import mesh_utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[CpuDevice(id=0) CpuDevice(id=1) CpuDevice(id=2) CpuDevice(id=3)]\n",
      " [CpuDevice(id=4) CpuDevice(id=5) CpuDevice(id=6) CpuDevice(id=7)]]\n",
      "Mesh('data': 2, 'model': 4)\n"
     ]
    }
   ],
   "source": [
    "# Create a mesh and annotate each axis with a name.\n",
    "device_mesh = mesh_utils.create_device_mesh((2, 4))\n",
    "print(device_mesh)\n",
    "\n",
    "mesh = Mesh(devices=device_mesh, axis_names=('data', 'model'))\n",
    "print(mesh)\n",
    "\n",
    "def mesh_sharding(pspec: PartitionSpec) -> NamedSharding:\n",
    "  return NamedSharding(mesh, pspec)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DotReluDot(nn.Module):\n",
    "  depth: int\n",
    "  dense_init: Callable = nn.initializers.xavier_normal()\n",
    "  @nn.compact\n",
    "  def __call__(self, x):\n",
    "\n",
    "    y = nn.Dense(self.depth,\n",
    "                 kernel_init=nn.with_partitioning(self.dense_init, (None, 'model')),\n",
    "                 use_bias=False,  # or overwrite with `bias_init`\n",
    "                 )(x)\n",
    "\n",
    "    y = jax.nn.relu(y)\n",
    "    # Force a local sharding annotation.\n",
    "    y = with_sharding_constraint(y, mesh_sharding(PartitionSpec('data', 'model')))\n",
    "\n",
    "    W2 = self.param(\n",
    "        'W2',\n",
    "        nn.with_partitioning(self.dense_init, ('model', None)),\n",
    "        (self.depth, x.shape[-1]))\n",
    "\n",
    "    z = jnp.dot(y, W2)\n",
    "    # Force a local sharding annotation.\n",
    "    z = with_sharding_constraint(z, mesh_sharding(PartitionSpec('data', None)))\n",
    "\n",
    "    # Return a tuple to conform with the API `flax.linen.scan` as shown in the cell below.\n",
    "    return z, None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MLP(nn.Module):\n",
    "  num_layers: int\n",
    "  depth: int\n",
    "  use_scan: bool\n",
    "  @nn.compact\n",
    "  def __call__(self, x):\n",
    "    if self.use_scan:\n",
    "      x, _ = nn.scan(DotReluDot, length=self.num_layers,\n",
    "                     variable_axes={\"params\": 0},\n",
    "                     split_rngs={\"params\": True},\n",
    "                     metadata_params={nn.PARTITION_NAME: None}\n",
    "                     )(self.depth)(x)\n",
    "    else:\n",
    "      for i in range(self.num_layers):\n",
    "        x, _ = DotReluDot(self.depth)(x)\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "# MLP hyperparameters.\n",
    "BATCH, LAYERS, DEPTH, USE_SCAN = 8, 4, 1024, False\n",
    "# Create fake inputs.\n",
    "x = jnp.ones((BATCH, DEPTH))\n",
    "# Initialize a PRNG key.\n",
    "k = random.key(0)\n",
    "\n",
    "# Create an Optax optimizer.\n",
    "optimizer = optax.adam(learning_rate=0.001)\n",
    "# Instantiate the model.\n",
    "model = MLP(LAYERS, DEPTH, USE_SCAN)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(8, 1024)"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                  CPU 0,1,2,3                                   </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                  CPU 4,5,6,7                                   </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m                                  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mCPU 0,1,2,3\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m                                   \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m                                  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mCPU 4,5,6,7\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m                                   \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "x_sharding = mesh_sharding(PartitionSpec('data', None)) # dimensions: (batch, length)\n",
    "x = jax.device_put(x, x_sharding)\n",
    "jax.debug.visualize_array_sharding(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "def init_fn(k, x, model, optimizer):\n",
    "  variables = model.init(k, x) # Initialize the model.\n",
    "  state = train_state.TrainState.create( # Create a `TrainState`.\n",
    "    apply_fn=model.apply,\n",
    "    params=variables['params'],\n",
    "    tx=optimizer)\n",
    "  return state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an abstract closure to wrap the function before feeding it in\n",
    "# because `jax.eval_shape` only takes pytrees as arguments.\n",
    "abstract_variables = jax.eval_shape(\n",
    "    functools.partial(init_fn, model=model, optimizer=optimizer), k, x)\n",
    "\n",
    "# This `state_sharding` has the same pytree structure as `state`, the output\n",
    "# of the `init_fn`.\n",
    "state_sharding = nn.get_sharding(abstract_variables, mesh)\n",
    "state_sharding"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "jax-env",
   "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}