{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "execute:\n",
    "  enabled: false\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=ndarray-md-1 fw=all -->\n",
    "\n",
    "# Data Manipulation\n",
    "\n",
    "In order to get anything done, \n",
    "we need some way to store and manipulate data.\n",
    "Generally, there are two important things \n",
    "we need to do with data: \n",
    "(i) acquire them; \n",
    "and (ii) process them once they are inside the computer. \n",
    "There is no point in acquiring data \n",
    "without some way to store it, \n",
    "so to start, let's get our hands dirty\n",
    "with $n$-dimensional arrays, \n",
    "which we also call *tensors*.\n",
    "If you already know the NumPy \n",
    "scientific computing package [@Harris.Millman.Walt.ea.2020], \n",
    "much of this will look familiar.\n",
    "For all modern deep learning frameworks,\n",
    "the *tensor class* (`ndarray` in MXNet, \n",
    "`Tensor` in PyTorch and TensorFlow,\n",
    "`jax.Array` in JAX) \n",
    "resembles NumPy's `ndarray`,\n",
    "with two additions.\n",
    "First, the tensor class\n",
    "supports automatic differentiation\n",
    "(that section).\n",
    "Second, it uses GPUs\n",
    "to accelerate numerical computation,\n",
    "whereas NumPy only runs on CPUs.\n",
    "These properties make neural networks\n",
    "both easy to code and fast to run.\n",
    "(We cover how to place tensors on a GPU and move them between devices\n",
    "in that section; until then, every tensor we create lives in\n",
    "CPU memory.)\n",
    "\n",
    "\n",
    "## Getting Started"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=ndarray-md-2 fw=jax -->\n",
    "\n",
    "To start, we import `jax` and its\n",
    "NumPy-like numerical interface `jax.numpy`,\n",
    "conventionally aliased as `jnp`."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "ndarray-getting-started-1"
   },
   "source": [
    "import jax\n",
    "from jax import numpy as jnp"
   ],
   "execution_count": 1,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=ndarray-md-3 fw=all -->\n",
    "\n",
    "A tensor represents a (possibly multidimensional) array of numerical values.\n",
    "In the one-dimensional case, i.e., when only one axis is needed for the data,\n",
    "a tensor is called a *vector*.\n",
    "With two axes, a tensor is called a *matrix*.\n",
    "With $k > 2$ axes, we drop the specialized names\n",
    "and just refer to the object as a $k^\\textrm{th}$-*order tensor*.\n",
    "\n",
    "Several functions create new tensors\n",
    "prepopulated with values.\n",
    "For example, by invoking `arange(n)`,\n",
    "we can create a vector of evenly spaced values,\n",
    "starting at 0 (included)\n",
    "and ending at `n` (not included).\n",
    "By default, the interval size is $1$.\n",
    "Unless otherwise specified,\n",
    "new tensors are stored in main memory\n",
    "and designated for CPU-based computation."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "ndarray-getting-started-2"
   },
   "source": [
    "x = jnp.arange(12, dtype=jnp.float32)\n",
    "x"
   ],
   "execution_count": 2,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "Array([ 0.,  1.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  9., 10., 11.],      dtype=float32)"
     },
     "metadata": {},
     "execution_count": 2
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=ndarray-md-4 fw=all -->\n",
    "\n",
    "The `dtype` (data type) of a tensor determines how its elements are stored.\n",
    "We use **32-bit** floating point (`float32`) as the default throughout these\n",
    "chapters. Training systems also use lower-precision floating-point formats,\n",
    "often together with higher-precision accumulation. that section\n",
    "develops those formats and their numerical tradeoffs. You can inspect a tensor's type via its\n",
    "`dtype` attribute. In finite precision, naive\n",
    "computations can overflow or underflow, which is why later sections compute\n",
    "quantities such as the softmax and cross-entropy with stabilized formulas\n",
    "(that section)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=ndarray-md-5 fw=jax -->\n",
    "\n",
    "Integer ranges like `arange` default to an integer type,\n",
    "which is why the cell above requests `dtype=jnp.float32` explicitly.\n",
    "To cast between types, use `astype`."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "ndarray-getting-started-3"
   },
   "source": [
    "x.size"
   ],
   "execution_count": 3,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "12"
     },
     "metadata": {},
     "execution_count": 3
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=ndarray-md-6 fw=all -->\n",
    "\n",
    "We can access a tensor's *shape* \n",
    "(the length along each axis)\n",
    "by inspecting its `shape` attribute.\n",
    "Because we are dealing with a vector here,\n",
    "the `shape` contains just a single element\n",
    "and is identical to the size."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "ndarray-getting-started-4"
   },
   "source": [
    "x.shape"
   ],
   "execution_count": 4,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(12,)"
     },
     "metadata": {},
     "execution_count": 4
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=ndarray-md-7 fw=all -->\n",
    "\n",
    "We can change the shape of a tensor\n",
    "without altering its size or values,\n",
    "by invoking `reshape`.\n",
    "For example, we can transform \n",
    "our vector `x` whose shape is (12,) \n",
    "to a matrix `X` with shape (3, 4).\n",
    "This new tensor retains all elements\n",
    "but reconfigures them into a matrix.\n",
    "Notice that the elements of our vector\n",
    "are laid out one row at a time and thus\n",
    "`x[3] == X[0, 3]`, as the figure illustrates.\n",
    "\n",
    "![Reshaping arranges the same 12 values into three rows without changing their order.](../img/ndarray-reshape.svg)"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "ndarray-getting-started-5"
   },
   "source": [
    "X = x.reshape(3, 4)\n",
    "X"
   ],
   "execution_count": 5,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "Array([[ 0.,  1.,  2.,  3.],\n       [ 4.,  5.,  6.,  7.],\n       [ 8.,  9., 10., 11.]], dtype=float32)"
     },
     "metadata": {},
     "execution_count": 5
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=ndarray-md-8 fw=all -->\n",
    "\n",
    "Note that specifying every shape component\n",
    "to `reshape` is redundant.\n",
    "Because we already know our tensor's size,\n",
    "we can work out one component of the shape given the rest.\n",
    "For example, given a tensor of size $n$\n",
    "and target shape ($h$, $w$),\n",
    "we know that $w = n/h$.\n",
    "To automatically infer one component of the shape,\n",
    "we can place a `-1` for the shape component\n",
    "that should be inferred automatically.\n",
    "In our case, instead of calling `x.reshape(3, 4)`,\n",
    "we could have equivalently called `x.reshape(-1, 4)` or `x.reshape(3, -1)`.\n",
    "\n",
    "For a contiguous tensor such as `x`, a framework can implement `reshape` by\n",
    "changing only shape metadata. More generally, `reshape` may return either a\n",
    "view that shares storage or a copy, depending on the layout and framework.\n",
    "Code should therefore not rely on storage sharing unless it explicitly asks\n",
    "for a view and satisfies that operation's layout requirements.\n",
    "\n",
    "Practitioners often need to work with tensors\n",
    "initialized to contain all 0s or 1s.\n",
    "We can construct a tensor with all elements set to 0 \n",
    "and a shape of (2, 3, 4) via the `zeros` function."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "ndarray-getting-started-6"
   },
   "source": [
    "jnp.zeros((2, 3, 4))"
   ],
   "execution_count": 6,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "Array([[[0., 0., 0., 0.],\n        [0., 0., 0., 0.],\n        [0., 0., 0., 0.]],\n\n       [[0., 0., 0., 0.],\n        [0., 0., 0., 0.],\n        [0., 0., 0., 0.]]], dtype=float32)"
     },
     "metadata": {},
     "execution_count": 6
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=ndarray-md-9 fw=all -->\n",
    "\n",
    "Similarly, we can create a tensor \n",
    "with all 1s by invoking `ones`."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "ndarray-getting-started-7"
   },
   "source": [
    "jnp.ones((2, 3, 4))"
   ],
   "execution_count": 7,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "Array([[[1., 1., 1., 1.],\n        [1., 1., 1., 1.],\n        [1., 1., 1., 1.]],\n\n       [[1., 1., 1., 1.],\n        [1., 1., 1., 1.],\n        [1., 1., 1., 1.]]], dtype=float32)"
     },
     "metadata": {},
     "execution_count": 7
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=ndarray-md-10 fw=all -->\n",
    "\n",
    "We often wish to \n",
    "sample each element randomly (and independently) \n",
    "from a given probability distribution.\n",
    "For example, the parameters of neural networks\n",
    "are often initialized randomly, in part to *break symmetry* between units:\n",
    "if every weight in a layer started at the same value (say, zero),\n",
    "the units would all compute the same thing and could never specialize.\n",
    "The following snippet creates a tensor \n",
    "with elements drawn from \n",
    "a standard Gaussian (normal) distribution\n",
    "with mean 0 and standard deviation 1."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "ndarray-getting-started-8"
   },
   "source": [
    "# Any call of a random function in JAX requires a key to be\n",
    "# specified, feeding the same key to a random function will\n",
    "# always result in the same sample being generated\n",
    "jax.random.normal(jax.random.key(0), (3, 4))"
   ],
   "execution_count": 8,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "Array([[ 1.6226422 ,  2.0252647 , -0.43359444, -0.07861735],\n       [ 0.1760909 , -0.97208923, -0.49529874,  0.4943786 ],\n       [ 0.6643493 , -0.9501635 ,  2.1795304 , -1.9551506 ]],      dtype=float32)"
     },
     "metadata": {},
     "execution_count": 8
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=ndarray-md-11 fw=all -->\n",
    "\n",
    "Random sampling raises the question of *reproducibility*. Calling a sampler\n",
    "repeatedly gives different draws, which is usually what we want, but for\n",
    "debugging or for figures that should not change between runs we fix a *seed*."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=ndarray-md-12 fw=jax -->\n",
    "\n",
    "JAX takes a more explicit stance: instead of a hidden global generator,\n",
    "every draw threads a *key* through the sampler, and the same key always\n",
    "yields the same sample, as the cell above shows."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=ndarray-md-13 fw=all -->\n",
    "\n",
    "Finally, we can construct tensors by\n",
    "supplying the exact values for each element\n",
    "via (possibly nested) Python list(s)\n",
    "containing numerical literals.\n",
    "Here, we construct a matrix with a list of lists,\n",
    "where the outermost list corresponds to axis 0,\n",
    "and the inner list corresponds to axis 1."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "ndarray-getting-started-9"
   },
   "source": [
    "jnp.array([[2, 1, 4, 3], [1, 2, 3, 4], [4, 3, 2, 1]])"
   ],
   "execution_count": 9,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "Array([[2, 1, 4, 3],\n       [1, 2, 3, 4],\n       [4, 3, 2, 1]], dtype=int32)"
     },
     "metadata": {},
     "execution_count": 9
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=ndarray-md-14 fw=all -->\n",
    "\n",
    "## Indexing and Slicing\n",
    "\n",
    "As with Python lists,\n",
    "we can access tensor elements \n",
    "by indexing (starting with 0).\n",
    "To access an element based on its position\n",
    "relative to the end of the list,\n",
    "we can use negative indexing.\n",
    "We can also access whole ranges of indices \n",
    "via slicing (e.g., `X[start:stop]`), \n",
    "where the returned value includes \n",
    "the first index (`start`) *but not the last* (`stop`).\n",
    "Finally, when only one index (or slice)\n",
    "is specified for a $k^\\textrm{th}$-order tensor,\n",
    "it is applied along axis 0.\n",
    "Thus, in the following code,\n",
    "`[-1]` selects the last row and `[1:3]`\n",
    "selects the second and third rows."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "ndarray-indexing-and-slicing-1"
   },
   "source": [
    "X[-1], X[1:3]"
   ],
   "execution_count": 10,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(Array([ 8.,  9., 10., 11.], dtype=float32),\n Array([[ 4.,  5.,  6.,  7.],\n        [ 8.,  9., 10., 11.]], dtype=float32))"
     },
     "metadata": {},
     "execution_count": 10
    }
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "ndarray-indexing-and-slicing-2"
   },
   "source": [
    "# JAX arrays are immutable. jax.numpy.ndarray.at index\n",
    "# update operators create a new array with the corresponding\n",
    "# modifications made\n",
    "X_new_1 = X.at[1, 2].set(17)\n",
    "X_new_1"
   ],
   "execution_count": 11,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "Array([[ 0.,  1.,  2.,  3.],\n       [ 4.,  5., 17.,  7.],\n       [ 8.,  9., 10., 11.]], dtype=float32)"
     },
     "metadata": {},
     "execution_count": 11
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=ndarray-md-15 fw=all -->\n",
    "\n",
    "If we want to assign multiple elements the same value,\n",
    "we apply the indexing on the left-hand side \n",
    "of the assignment operation.\n",
    "For instance, `[:2, :]` accesses \n",
    "the first and second rows,\n",
    "where `:` takes all the elements along axis 1 (column).\n",
    "While we discussed indexing for matrices,\n",
    "this also works for vectors\n",
    "and for tensors of more than two dimensions."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "ndarray-indexing-and-slicing-3"
   },
   "source": [
    "X_new_2 = X_new_1.at[:2, :].set(12)\n",
    "X_new_2"
   ],
   "execution_count": 12,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "Array([[12., 12., 12., 12.],\n       [12., 12., 12., 12.],\n       [ 8.,  9., 10., 11.]], dtype=float32)"
     },
     "metadata": {},
     "execution_count": 12
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=ndarray-md-16 fw=all -->\n",
    "\n",
    "## Operations\n",
    "\n",
    "Now that we know how to construct tensors\n",
    "and how to read from and write to their elements,\n",
    "we can begin to manipulate them\n",
    "with various mathematical operations.\n",
    "Among the most useful of these \n",
    "are the *elementwise* operations.\n",
    "These apply a standard scalar operation\n",
    "to each element of a tensor.\n",
    "For functions that take two tensors as inputs,\n",
    "elementwise operations apply some standard binary operator\n",
    "on each pair of corresponding elements.\n",
    "We can create an elementwise function \n",
    "from any function that maps \n",
    "from a scalar to a scalar.\n",
    "\n",
    "In mathematical notation, we denote such\n",
    "*unary* scalar operators (taking one input)\n",
    "by the signature \n",
    "$f: \\mathbb{R} \\rightarrow \\mathbb{R}$.\n",
    "This just means that the function maps\n",
    "from any real number onto some other real number.\n",
    "Most standard operators, including unary ones like $e^x$, can be applied elementwise."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "ndarray-operations-1"
   },
   "source": [
    "jnp.exp(x)"
   ],
   "execution_count": 13,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "Array([1.0000000e+00, 2.7182817e+00, 7.3890562e+00, 2.0085537e+01,\n       5.4598148e+01, 1.4841316e+02, 4.0342880e+02, 1.0966332e+03,\n       2.9809580e+03, 8.1030840e+03, 2.2026467e+04, 5.9874145e+04],      dtype=float32)"
     },
     "metadata": {},
     "execution_count": 13
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=ndarray-md-17 fw=all -->\n",
    "\n",
    "Likewise, we denote *binary* scalar operators,\n",
    "which map pairs of real numbers\n",
    "to a (single) real number\n",
    "via the signature \n",
    "$f: \\mathbb{R}, \\mathbb{R} \\rightarrow \\mathbb{R}$.\n",
    "Given any two vectors $\\mathbf{u}$ \n",
    "and $\\mathbf{v}$ *of the same shape*,\n",
    "and a binary operator $f$, we can produce a vector\n",
    "$\\mathbf{c} = F(\\mathbf{u},\\mathbf{v})$\n",
    "by setting $c_i \\gets f(u_i, v_i)$ for all $i$,\n",
    "where $c_i, u_i$, and $v_i$ are the $i^\\textrm{th}$ elements\n",
    "of vectors $\\mathbf{c}, \\mathbf{u}$, and $\\mathbf{v}$.\n",
    "Here, we produced the vector-valued\n",
    "$F: \\mathbb{R}^d, \\mathbb{R}^d \\rightarrow \\mathbb{R}^d$\n",
    "by *lifting* the scalar function\n",
    "to an elementwise vector operation.\n",
    "The common standard arithmetic operators\n",
    "for addition (`+`), subtraction (`-`), \n",
    "multiplication (`*`), division (`/`), \n",
    "and exponentiation (`**`)\n",
    "have all been *lifted* to elementwise operations\n",
    "for identically-shaped tensors of arbitrary shape."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "ndarray-operations-2"
   },
   "source": [
    "x = jnp.array([1.0, 2, 4, 8])\n",
    "y = jnp.array([2, 2, 2, 2])\n",
    "x + y, x - y, x * y, x / y, x ** y"
   ],
   "execution_count": 14,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(Array([ 3.,  4.,  6., 10.], dtype=float32),\n Array([-1.,  0.,  2.,  6.], dtype=float32),\n Array([ 2.,  4.,  8., 16.], dtype=float32),\n Array([0.5, 1. , 2. , 4. ], dtype=float32),\n Array([ 1.,  4., 16., 64.], dtype=float32))"
     },
     "metadata": {},
     "execution_count": 14
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=ndarray-md-18 fw=all -->\n",
    "\n",
    "In addition to elementwise computations,\n",
    "we can also perform linear algebraic operations,\n",
    "such as dot products and matrix multiplications.\n",
    "We will elaborate on these\n",
    "in that section.\n",
    "\n",
    "We can also *concatenate* multiple tensors,\n",
    "stacking them end-to-end to form a larger one.\n",
    "We just need to provide a list of tensors\n",
    "and tell the system along which axis to concatenate.\n",
    "The example below shows what happens when we concatenate\n",
    "two matrices along rows (axis 0)\n",
    "instead of columns (axis 1).\n",
    "We can see that the first output's axis-0 length ($6$)\n",
    "is the sum of the two input tensors' axis-0 lengths ($3 + 3$);\n",
    "while the second output's axis-1 length ($8$)\n",
    "is the sum of the two input tensors' axis-1 lengths ($4 + 4$)."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "ndarray-operations-3"
   },
   "source": [
    "X = jnp.arange(12, dtype=jnp.float32).reshape((3, 4))\n",
    "Y = jnp.array([[2.0, 1, 4, 3], [1, 2, 3, 4], [4, 3, 2, 1]])\n",
    "jnp.concatenate((X, Y), axis=0), jnp.concatenate((X, Y), axis=1)"
   ],
   "execution_count": 15,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(Array([[ 0.,  1.,  2.,  3.],\n        [ 4.,  5.,  6.,  7.],\n        [ 8.,  9., 10., 11.],\n        [ 2.,  1.,  4.,  3.],\n        [ 1.,  2.,  3.,  4.],\n        [ 4.,  3.,  2.,  1.]], dtype=float32),\n Array([[ 0.,  1.,  2.,  3.,  2.,  1.,  4.,  3.],\n        [ 4.,  5.,  6.,  7.,  1.,  2.,  3.,  4.],\n        [ 8.,  9., 10., 11.,  4.,  3.,  2.,  1.]], dtype=float32))"
     },
     "metadata": {},
     "execution_count": 15
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=ndarray-md-19 fw=all -->\n",
    "\n",
    "Sometimes, we want to \n",
    "construct a boolean tensor via *logical statements*.\n",
    "Take `X == Y` as an example.\n",
    "For each position `i, j`, if `X[i, j]` and `Y[i, j]` are equal, \n",
    "then the corresponding entry in the result is `True`,\n",
    "otherwise it is `False`.\n",
    "(Booleans behave as 1 and 0 in arithmetic, so summing\n",
    "such a mask counts the positions where the two tensors agree.)"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "ndarray-operations-4"
   },
   "source": [
    "X == Y"
   ],
   "execution_count": 16,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "Array([[False,  True, False,  True],\n       [False, False, False, False],\n       [False, False, False, False]], dtype=bool)"
     },
     "metadata": {},
     "execution_count": 16
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=ndarray-md-20 fw=all -->\n",
    "\n",
    "Summing all the elements in the tensor yields a tensor with only one element."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "ndarray-operations-5"
   },
   "source": [
    "X.sum()"
   ],
   "execution_count": 17,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "Array(66., dtype=float32)"
     },
     "metadata": {},
     "execution_count": 17
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=ndarray-md-21 fw=all -->\n",
    "\n",
    "## Broadcasting\n",
    "\n",
    "By now, you know how to perform \n",
    "elementwise binary operations\n",
    "on two tensors of the same shape. \n",
    "Under certain conditions,\n",
    "even when shapes differ, \n",
    "we can still perform elementwise binary operations\n",
    "by invoking the *broadcasting mechanism*\n",
    "(a convention inherited from NumPy).\n",
    "\n",
    "Broadcasting follows a simple rule. Line the two shapes up\n",
    "from the *right* and compare them axis by axis: two axes are compatible when\n",
    "they are equal or when one of them is $1$ (a missing leading axis counts as\n",
    "$1$). The size-$1$ axis is then *stretched* to match the other, by virtually\n",
    "copying its entries, and the elementwise operation is applied to the\n",
    "resulting arrays. If any pair\n",
    "of axes is incompatible, the operation raises an error rather than guessing.\n",
    "the figure shows the mechanism at work on the example\n",
    "that follows: each size-$1$ axis is stretched until both operands share the\n",
    "shape $3\\times2$.\n",
    "\n",
    "![Broadcasting stretches the size-1 axes: a $3\\times1$ column and a $1\\times2$ row each expand to $3\\times2$ before the elementwise addition.](../img/ndarray-broadcasting.svg)"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "ndarray-broadcasting-1"
   },
   "source": [
    "a = jnp.arange(3).reshape((3, 1))\n",
    "b = jnp.arange(2).reshape((1, 2))\n",
    "a, b"
   ],
   "execution_count": 18,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(Array([[0],\n        [1],\n        [2]], dtype=int32),\n Array([[0, 1]], dtype=int32))"
     },
     "metadata": {},
     "execution_count": 18
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=ndarray-md-22 fw=all -->\n",
    "\n",
    "Since `a` and `b` are $3\\times1$ \n",
    "and $1\\times2$ matrices, respectively,\n",
    "their shapes do not match up.\n",
    "Broadcasting produces a larger $3\\times2$ matrix \n",
    "by replicating matrix `a` along the columns\n",
    "and matrix `b` along the rows\n",
    "before adding them elementwise."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "ndarray-broadcasting-2"
   },
   "source": [
    "a + b"
   ],
   "execution_count": 19,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "Array([[0, 1],\n       [1, 2],\n       [2, 3]], dtype=int32)"
     },
     "metadata": {},
     "execution_count": 19
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=ndarray-md-23 fw=all -->\n",
    "\n",
    "What happens when the shapes are *not* compatible? Lining up $(3, 2)$ and\n",
    "$(2, 3)$ from the right pairs $2$ with $3$ and $3$ with $2$: neither axis pair\n",
    "matches and neither member is $1$, so the framework refuses rather than\n",
    "guessing."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "ndarray-broadcasting-3"
   },
   "source": [
    "try:\n",
    "    jnp.ones((3, 2)) + jnp.ones((2, 3))\n",
    "except Exception as e:\n",
    "    print(e)"
   ],
   "execution_count": 20,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "add got incompatible shapes for broadcasting: (3, 2), (2, 3).\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=ndarray-md-24 fw=all -->\n",
    "\n",
    "## Saving Memory\n",
    "\n",
    "Running operations can cause new memory to be\n",
    "allocated to host results.\n",
    "For example, if we write `Y = Y + X`,\n",
    "we dereference the tensor that `Y` used to point to\n",
    "and instead point `Y` at the newly allocated memory.\n",
    "We can demonstrate the rebinding with Python's `id()` function,\n",
    "which identifies a Python object for the duration of its lifetime.\n",
    "Note that after we run `Y = Y + X`,\n",
    "`id(Y)` changes.\n",
    "That is because Python first evaluates `Y + X`,\n",
    "allocating new memory for the result \n",
    "and then binds `Y` to the new tensor object. This test says nothing about the\n",
    "address of the tensor's underlying storage."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "ndarray-saving-memory-1"
   },
   "source": [
    "before = id(Y)\n",
    "Y = Y + X\n",
    "id(Y) == before"
   ],
   "execution_count": 21,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "False"
     },
     "metadata": {},
     "execution_count": 21
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=ndarray-md-25 fw=all -->\n",
    "\n",
    "This might be undesirable for two reasons.\n",
    "First, we do not want to run around\n",
    "allocating memory unnecessarily all the time.\n",
    "In machine learning, we often have\n",
    "gigabytes of parameters\n",
    "and update all of them multiple times per second.\n",
    "In frameworks that support mutation, in-place updates can avoid these\n",
    "allocations. They must be used with care because automatic differentiation\n",
    "may need an earlier value to compute a gradient; functional and compiled\n",
    "systems may instead reuse buffers internally without exposing mutation.\n",
    "Second, we might point at the \n",
    "same parameters from multiple variables.\n",
    "If we do not update in place, \n",
    "we must be careful to update all of these references,\n",
    "lest we spring a memory leak \n",
    "or inadvertently refer to stale parameters."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "ndarray-saving-memory-2"
   },
   "source": [
    "# JAX arrays do not allow in-place operations"
   ],
   "execution_count": 22,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "ndarray-saving-memory-3"
   },
   "source": [
    "# JAX arrays are immutable, so functional updates return new arrays.\n",
    "# Use the `.at[...].set(...)` syntax; under JIT, XLA can fuse these\n",
    "# updates and reuse buffers, recovering most of the in-place benefit.\n",
    "X_new = X.at[:].set(X + Y)\n",
    "id(X_new) == id(X)"
   ],
   "execution_count": 23,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "False"
     },
     "metadata": {},
     "execution_count": 23
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=ndarray-md-26 fw=all -->\n",
    "\n",
    "## Conversion to Other Python Objects"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=ndarray-md-27 fw=jax -->\n",
    "\n",
    "Converting to a NumPy tensor (`ndarray`), or vice versa, is easy:\n",
    "`jax.device_get` copies a JAX array to the host as a NumPy array,\n",
    "and `jax.device_put` transfers it back.\n",
    "The converted result does not share memory,\n",
    "since a JAX array may live on an accelerator."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "ndarray-conversion-to-other-python-objects-1"
   },
   "source": [
    "A = jax.device_get(X)\n",
    "B = jax.device_put(A)\n",
    "type(A), type(B)"
   ],
   "execution_count": 24,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(numpy.ndarray, jaxlib._jax.ArrayImpl)"
     },
     "metadata": {},
     "execution_count": 24
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=ndarray-md-28 fw=all -->\n",
    "\n",
    "To convert a size-1 tensor to a Python scalar,\n",
    "we can invoke the `item` method or Python's built-in `float` and `int`."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "ndarray-conversion-to-other-python-objects-2"
   },
   "source": [
    "a = jnp.array(3.5)\n",
    "a, a.item(), float(a), int(a)"
   ],
   "execution_count": 25,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(Array(3.5, dtype=float32, weak_type=True), 3.5, 3.5, 3)"
     },
     "metadata": {},
     "execution_count": 25
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=ndarray-md-29 fw=all -->\n",
    "\n",
    "## Discussion\n",
    "\n",
    "The tensor class is the main interface for storing and manipulating data in deep learning libraries.\n",
    "Tensors provide a variety of functionalities including construction routines; indexing and slicing; basic mathematics operations; broadcasting; memory-efficient assignment; and conversion to and from other Python objects.\n",
    "\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. Run the code in this section. Change the conditional statement `X == Y` to `X < Y` or `X > Y`, and then see what kind of tensor you can get.\n",
    "1. Replace the two tensors that operate by element in the broadcasting mechanism with other shapes, e.g., 3-dimensional tensors. Is the result the same as expected?\n",
    "1. Create the vector `x = arange(24)` and reshape it into a $(2, 3, 4)$ tensor using `-1` for one of the components. Which component does the framework infer, and why is at most one `-1` allowed?\n",
    "1. Build a $(3, 4)$ tensor and predict, *before running*, the shape of its sum along `axis=0`, along `axis=1`, and with `keepdims=True`. Then check each prediction against the code.\n",
    "1. Try to add two tensors whose shapes are *not* broadcast-compatible, for example shapes $(3, 2)$ and $(2, 3)$. What error do you get? Use the alignment rule from the broadcasting section to explain it, then find a reshape of one operand that makes the addition valid.\n",
    "1. Recall the saving-memory discussion. Use `id()` to verify that `X[:] = X + Y` (or `X += Y`) keeps the same Python tensor object while `X = X + Y` binds `X` to a new object. Why does object identity alone not prove that the underlying storage was reused? Consult your framework's storage or buffer API and test storage identity where such an API exists."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=ndarray-md-30 fw=jax -->\n",
    "\n",
    "[Discussions](https://d2l.discourse.group/t/17966)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-jax",
   "display_name": "d2l (JAX)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
