torchdim.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bbbc39f0"
      },
      "source": [
        "Named Tensors using First-class Dimensions in PyTorch\n",
        "=====================================================\n",
        "\n",
        "An implementation of [named tensors](https://namedtensor.github.io) with the functionality of [einsum](http://einops.rocks]http://einops.rocks) , batching ([vmap](https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap), [xmap](https://jax.readthedocs.io/en/latest/notebooks/xmap_tutorial.html)), and tensor indexing by adding dimension objects to PyTorch_.\n",
        "\n",
        "The tensor input to a resnet might have the shape [8, 3, 224, 224] but informally we think of those dimensions as 'batch', 'channel', 'width', and 'height'. Eventhough 'width' and 'height' have the same _size_ we still think of them as separate dimensions, and if we have two _different_ images, we think of both as sharing the _same_ 'channel' dimension.\n",
        "\n",
        " Instead of treating this concept informally, named tensors gives these dimensions names. Instead of using strings to name dimensions like [PyTorch's current implementation](https://pytorch.org/docs/stable/named_tensor.html), this library introduces a Python object, a `Dim`, to represent the concept. By expanding the semantics of tensors with dim objects, in addition to naming dimensions, we can get behavior equivalent to batching transforms (xmap, vmap), einops-style rearragement, and loop-style tensor indexing.\n",
        "\n",
        "\n",
        "    # einsum\n",
        "    def mm(A, B):\n",
        "      i, j, k = dims(3)\n",
        "      r = (A[i, k] * B[k, j]).sum(k)\n",
        "      return r.order(i, j)\n",
        "    \n",
        "    # rearrange\n",
        "    def pixel_shuffle(img, upscale_factor=2):\n",
        "      h2, w2, c, b, h, w = dims(6)\n",
        "      h2.size = w2.size = upscale_factor\n",
        "      return img[b, (c, h2, w2), h, w].order(b, c, (h, h2), (w, w2))\n",
        "    \n",
        "    # batching\n",
        "    def bmm(A, B):\n",
        "      i = dims(1)\n",
        "      return mm(A[i], B[i]).order(i)\n",
        "\n",
        "    # indexing\n",
        "    def embedding_bag(input, embedding_weights):\n",
        "      batch, sequence, features = dims(3)\n",
        "      r = embedding_weights[input[batch, sequence], features].sum(sequence)\n",
        "      return r.order(batch, features)\n",
        "    \n",
        "\n"
      ],
      "id": "bbbc39f0"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lk0qjXoEWVrc"
      },
      "source": [
        "Colab Setup\n",
        "-----------"
      ],
      "id": "lk0qjXoEWVrc"
    },
    {
      "cell_type": "code",
      "execution_count": 44,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "sIrEZVg6yBXF",
        "outputId": "fabff575-0da7-43c6-b7a3-193ae90e384e"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting torch==1.13.0.dev20220614+cpu\n",
            "  Using cached https://download.pytorch.org/whl/nightly/cpu/torch-1.13.0.dev20220614%2Bcpu-cp37-cp37m-linux_x86_64.whl (190.8 MB)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from torch==1.13.0.dev20220614+cpu) (4.1.1)\n"
          ]
        }
      ],
      "source": [
        "!pip3 install  https://download.pytorch.org/whl/nightly/cpu/torch-1.13.0.dev20220614%2Bcpu-cp37-cp37m-linux_x86_64.whl -U\n",
        "# if you want to try CUDA:\n",
        "# !pip3 install --pre torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/cu113 -U\n"
      ],
      "id": "sIrEZVg6yBXF"
    },
    {
      "cell_type": "code",
      "execution_count": 45,
      "metadata": {
        "id": "vExaBUzb2aa4",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "bb8e7089-7f9b-4a5b-c90b-40d51ec09c37"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting torchdim==1.0\n",
            "  Downloading https://www.dropbox.com/s/e4rgdqtirvblmnl/torchdim-1.0-cp37-cp37m-linux_x86_64.whl?dl=1 (20.8 MB)\n",
            "\u001b[K     |████████████████████████████████| 20.8 MB 2.3 MB/s \n",
            "\u001b[?25h"
          ]
        }
      ],
      "source": [
        "!pip install https://www.dropbox.com/s/e4rgdqtirvblmnl/torchdim-1.0-cp37-cp37m-linux_x86_64.whl?dl=1"
      ],
      "id": "vExaBUzb2aa4"
    },
    {
      "cell_type": "code",
      "source": [
        "# from getpass import getpass\n",
        "# username = input('Github username')\n",
        "# oauth_key = getpass('oauth access token')\n",
        "# !pip install ninja  # Makes the build go faster\n",
        "# !git clone https://{username}:{oauth_key}@github.com/fairinternal/torchdim\n",
        "# !(cd torchdim; git submodule update --init --recursive; python setup.py bdist_wheel)\n",
        "# !pip install torchdim/dist/torchdim-1.0-cp37-cp37m-linux_x86_64.whl"
      ],
      "metadata": {
        "id": "jm_rjn6h-4Bz"
      },
      "id": "jm_rjn6h-4Bz",
      "execution_count": 46,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "x5NeLkrbyBxl"
      },
      "source": [
        "Creating and Binding Dims\n",
        "=========================\n",
        "\n",
        "Python objects that represent dimension are created using the `dims` operator.[^1]"
      ],
      "id": "x5NeLkrbyBxl"
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "3f2a8718"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "from torchdim import dims\n",
        "\n",
        "batch, channel, width, height = dims(4)"
      ],
      "id": "3f2a8718"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4607f4c5"
      },
      "source": [
        "The existing implemention of [Named Tensors](https://pytorch.org/docs/stable/named_tensor.html) in PyTorch, or [JAX's xmap](https://jax.readthedocs.io/en/latest/notebooks/xmap_tutorial.html) use strings to name dimensions. We call these dimensions _first class_ because they are Python objects.\n",
        "\n",
        "In addition to the normal _positional_ dimensions in a tensor, tensors can also have a separate set of first-class dimensions.\n",
        "\n",
        "You can create tensors with first-class dimensions by indexing the normal positional dimensions of a tensor with a dimension object. The `ndim` property continues to list the number of positional dimesions, while the new `dims` property lists all the bound first-class dimensions."
      ],
      "id": "4607f4c5"
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "53a55c84",
        "outputId": "58e77535-5f6c-457c-8af4-184c48f10ea6"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "4\n"
          ]
        }
      ],
      "source": [
        "input = torch.rand(2, 3, 224, 224)\n",
        "print(input.ndim)"
      ],
      "id": "53a55c84"
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8511a637",
        "outputId": "63956be3-dbbb-4161-83a5-8fe4185bca9d"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "(batch, channel, width, height)\n"
          ]
        }
      ],
      "source": [
        "input_fc = input[batch, channel, width, height]\n",
        "print(input_fc.dims) # first class dimensions"
      ],
      "id": "8511a637"
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "af2d9bdf",
        "outputId": "8a301824-f8ef-43a2-da8a-34441964a964"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "0\n"
          ]
        }
      ],
      "source": [
        "\n",
        "# since we converted all the positional dimesions\n",
        "# first class `input_fc` has 0 positional dimensions now.\n",
        "print(input_fc.ndim) "
      ],
      "id": "af2d9bdf"
    },
    {
      "cell_type": "markdown",
      "source": [
        "Notice that indexing creates a _new_ Tensor, `input_fc` with bound first-class dimensions. It does not modify the original tensor `input`, which still has 4 positional dimensions."
      ],
      "metadata": {
        "id": "k_QU4SAgV7S1"
      },
      "id": "k_QU4SAgV7S1"
    },
    {
      "cell_type": "code",
      "source": [
        "print(input.ndim) # unchanged  "
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4L3RVV-pWOwe",
        "outputId": "5e3bb2d0-4b7d-4438-f412-f7eebf75a855"
      },
      "id": "4L3RVV-pWOwe",
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "4\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Importantly, indexing with square brackets _applies only to positional dimensions_, so attempting to index a tensor with only first class dims will error[^2]:"
      ],
      "metadata": {
        "id": "ATCGfUuMXgfW"
      },
      "id": "ATCGfUuMXgfW"
    },
    {
      "cell_type": "code",
      "source": [
        "try:\n",
        "  input_fc[0]\n",
        "except ValueError as ve:\n",
        "  print(ve)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "IgubYKeHX8QM",
        "outputId": "c9d05459-acb2-4298-a091-3bda885f1cef"
      },
      "id": "IgubYKeHX8QM",
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "at least 1 indices were supplied but the tensor only has 0 dimensions\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Generally, it is possible to construct tensors with a mixture of positional and first class dimensions:"
      ],
      "metadata": {
        "id": "0tJWQvfFWNrV"
      },
      "id": "0tJWQvfFWNrV"
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "f8a63a88",
        "outputId": "7b6d9ae7-908c-4275-d004-0466d596631f"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "(batch, height)\n"
          ]
        }
      ],
      "source": [
        "input_mixed = input[batch, :, :, height]\n",
        "print(input_mixed.dims)"
      ],
      "id": "f8a63a88"
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "84a28b41",
        "outputId": "598029c5-f888-4045-dcca-63a6080757df"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "2\n"
          ]
        }
      ],
      "source": [
        "print(input_mixed.ndim)"
      ],
      "id": "84a28b41"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "99d3e0eb"
      },
      "source": [
        "Dimension Sizes\n",
        "---------------\n",
        "\n",
        "Dimensions will take on the size of the first thing they are bound to:"
      ],
      "id": "99d3e0eb"
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "33a71c9d",
        "outputId": "0c0ce2e0-721e-4159-9799-5f045bad85e3"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "3\n"
          ]
        }
      ],
      "source": [
        "input = torch.rand(3)\n",
        "x = dims(1)\n",
        "input_fc = input[x]\n",
        "print(x.size)"
      ],
      "id": "33a71c9d"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2edc5f2a"
      },
      "source": [
        "But you can also directly set the size of dimension:"
      ],
      "id": "2edc5f2a"
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "id": "f52675e2",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "6d9880c2-e09a-4cd2-c12f-65e88072e4b8"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Dim 'i' previously bound to a dimension of size 5 cannot bind to a dimension of size 3\n"
          ]
        }
      ],
      "source": [
        "i = dims(1)\n",
        "\n",
        "i.size = 5 # ok, i previously did not have a size\n",
        "\n",
        "i.size = 5 # ok, it already had the size 5\n",
        "try:\n",
        "    i.size = 3\n",
        "except Exception as e:\n",
        "  print(e)"
      ],
      "id": "f52675e2"
    },
    {
      "cell_type": "code",
      "source": [
        "j = dims(sizes=[4]) # can also be set on construction"
      ],
      "metadata": {
        "id": "natJxrX14_WT"
      },
      "id": "natJxrX14_WT",
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2d86e49b"
      },
      "source": [
        "[^1]: We use a bit of Python introspection to set the debug names for the dimensions based on the names of the variables they are assigned to.\n",
        "[^2]: Indexing of first-class dimensions can be done with the `index` method by specifying the dimension to be index into (e.g. `input_fc.index(batch, 0)`.\n",
        "\n",
        "Semantics of Dimensions\n",
        "=======================\n",
        "The power of named tensors arises from how the first-class dimensions in the Tensors composed with existing operations.\n",
        "\n",
        "Three rules define how dimension objects behave with existing Tensors. \n",
        "\n",
        "Rule 1: Implicit Batching\n",
        "-------------------------\n",
        "**Tensor operations (e.g. `input + bias`) are implicitly batched over the union of the first-class dimensions in their inputs.**\n",
        "\n",
        "If `input` has dimensions `batch, channel` and `bias` has dimension `channel`, the output will have the union of those dimensions (`batch, channel`), and the result will computed as if there was a loop over all the first-class dimensions.[^3]"
      ],
      "id": "2d86e49b"
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "58ae274b",
        "outputId": "724cd1e1-94f8-46f8-f91e-b8f88e69f2d2"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "(batch, channel)\n"
          ]
        }
      ],
      "source": [
        "input_positional = torch.rand(128, 32)\n",
        "bias_positional = torch.rand(32)\n",
        "\n",
        "batch, channel = dims(2)\n",
        "input = input_positional[batch, channel]\n",
        "bias = bias_positional[channel]\n",
        "\n",
        "result = input + bias\n",
        "print(result.dims)"
      ],
      "id": "58ae274b"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "89f5bc44"
      },
      "source": [
        "It is helpful think of operators on tensors with first-class dimensions by analogy to code with explicit loops over dimensions, with the first-class dimensions of the inputs acting as implicit `for` loops, and the values in the tensor being scalars within the body of the loop:\n",
        "\n",
        "    # mental model: loop-level analogy\n",
        "    for batch in range(batch.size):\n",
        "        for channel in range(channel.size):\n",
        "           input = input_positional[batch, channels]\n",
        "           bias = bias_positional[channels] \n",
        "           result[batch, channels] =  input + bias # arithmetic on scalars\n",
        "\n",
        "Positional dimensions behave as they did before (e.g. for + they will broadcast), and can be thought of as being a standard tensor _used within the implicit loops_ defined by first-class dimensions.\n",
        "\n",
        "In this example, we broke down the expression into lines that bind the dimension to positional tensors and then another line to do the compute. In practice, we often combine these in one statement:\n"
      ],
      "id": "89f5bc44"
    },
    {
      "cell_type": "code",
      "source": [
        "result = input_positional[batch, channel] + bias_positional[channel]\n",
        "result.dims"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Z5wk9VChcsVI",
        "outputId": "385bc5f5-783f-4c09-ff6f-cd94095b505f"
      },
      "id": "Z5wk9VChcsVI",
      "execution_count": 16,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(batch, channel)"
            ]
          },
          "metadata": {},
          "execution_count": 16
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "[^3] This rule is similar to how named dimensions in xmap behave within a function, but instead of introducing the dimensions via a functional transform, they are bound on the objects using indexing.\n"
      ],
      "metadata": {
        "id": "VmpFzi2XczWB"
      },
      "id": "VmpFzi2XczWB"
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "Rule 2: Specifying dimensions\n",
        "-----------------------------\n",
        "**Wherever an integer is used to specify a dimension in the existing torch operator, a first-class dimensions can be used instead to tell the operator to work over that dimension.**"
      ],
      "metadata": {
        "id": "CLGkBY6nct0U"
      },
      "id": "CLGkBY6nct0U"
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7c315f98",
        "outputId": "d86d9709-7caa-4dd4-bbc3-57c4116340b4"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "(batch, channel)\n"
          ]
        }
      ],
      "source": [
        "batch, channel, width, height = dims(4)\n",
        "input_positional = torch.rand(2, 3, 224, 224)\n",
        "input = input_positional[batch, channel, width, height]\n",
        "avg_pixel_color = input.mean((width, height))\n",
        "\n",
        "print(avg_pixel_color.dims)"
      ],
      "id": "7c315f98"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d8ace9bb"
      },
      "source": [
        "Any other first-class dimensions (e.g. batch, channel) are still implicitly batched according to Rule #1.\n",
        "\n",
        "Rule 3: Dims are Tensors\n",
        "------------------------\n",
        "**A first-class dimension `d` can be used wherever a Tensor is expected. It will act as if it were a tensor whose only dimension is itself, `d`, and the values along the dimension are the indices of each entry `(0, 1, 2, ..., d.size - 1)`**"
      ],
      "id": "d8ace9bb"
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "bc239aab",
        "outputId": "05e67c7c-5e1c-40f7-ac12-1d1267136b06"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "(channel,)\n"
          ]
        }
      ],
      "source": [
        "print(channel.dims)"
      ],
      "id": "bc239aab"
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "64a99d2f",
        "outputId": "1463f6bc-3e1d-48e2-f204-136f88988247"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tensor([1000, 1001, 1002])\n",
            "with dims=(channel,) sizes=(3,)\n"
          ]
        }
      ],
      "source": [
        "print(channel + 1000)"
      ],
      "id": "64a99d2f"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "496e7877"
      },
      "source": [
        "This means that a dimension used as a tensor acts as an index into that dimension. Going back to our loop-level analogy, it is analogous to using the loop variable as a value:\n",
        "\n",
        "    # mental model: loop-level analogy\n",
        "    for channels in range(batch.size):\n",
        "        result[channels] = channels + 1000\n"
      ],
      "id": "496e7877"
    },
    {
      "cell_type": "markdown",
      "source": [
        "Arithmetic using dimension indices comes up a lot, such as the mask for an upper triangular part of a matrix. Using dims as tensors makes it easy:"
      ],
      "metadata": {
        "id": "xVty_z5obX9B"
      },
      "id": "xVty_z5obX9B"
    },
    {
      "cell_type": "code",
      "source": [
        "from torchdim import dims\n",
        "i, j = dims(sizes=[4, 4])\n",
        "print(i <= j)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "PIZzBPyHbtdw",
        "outputId": "4dedf8ad-16ce-4c59-e45a-04e37629892e"
      },
      "id": "PIZzBPyHbtdw",
      "execution_count": 20,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tensor([[ True,  True,  True,  True],\n",
            "        [False,  True,  True,  True],\n",
            "        [False, False,  True,  True],\n",
            "        [False, False, False,  True]])\n",
            "with dims=(i, j) sizes=(4, 4)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "Because of the intentional similarity to loop-level code, using dimsions as tensors makes complicated indexing arithmetic easier to read.\n",
        "\n",
        "Here is code that lookups up features in an embedding table given a sequence of ids:"
      ],
      "metadata": {
        "id": "chlxDtSqbW36"
      },
      "id": "chlxDtSqbW36"
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "d4a2fe7d",
        "outputId": "dce94c0c-b94e-4013-a994-bd0de3e78ac5"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "(sequence, features)\n"
          ]
        }
      ],
      "source": [
        "sequence, features = dims(2)\n",
        "embeddings = torch.rand(8, 128)\n",
        "words = torch.tensor([5, 4, 0,])\n",
        "\n",
        "state = embeddings[words[sequence], features]\n",
        "print(state.dims)"
      ],
      "id": "d4a2fe7d"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7ba139ce"
      },
      "source": [
        "With the following analogy to loops:\n",
        "\n",
        "    # mental model: loop-level analogy\n",
        "\n",
        "    for sequence in range(words.size(0)):\n",
        "        for features in range(embeddings.size(1)):\n",
        "            state = embeddings[words[sequence], features]\n",
        "\n",
        "Earlier we showed how binding tensors dimension is done with indexing `A[i, j]`. In fact, this binding is just the normal indexing operator. Its behavior follows directly from the behavior of indexing with tensor indices combined with Rule #3 and Rule #1. The expression `A[i + 1, j]` also creates a tensor with dimensions `i` and `j` but with different indexing math. The implementation knows when simple indexing patterns are used and only actually runs a kernel to do indexing when needed.\n",
        "\n",
        "Unbinding Dims\n",
        "-------------\n",
        "The `order` method converts first-class dimensions in a tensor back to normal positional dimensions by specifying an order for those dimensions.[^4]\n",
        "\n",
        "By specifiying a different order from how things were originally bound, it is easy to do transpositions."
      ],
      "id": "7ba139ce"
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {
        "id": "23378aea"
      },
      "outputs": [],
      "source": [
        "i, j = dims(2)\n",
        "A = torch.rand(3, 4)\n",
        "A_T = A[i, j].order(j, i)\n",
        "assert torch.allclose(A.T, A_T)"
      ],
      "id": "23378aea"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jdPvixK_xsm5"
      },
      "source": [
        "Indexing acts left-to-right, and `order` also places the new dimensions back on the left, so it possible to work on tensors that have mixed positonal and first-class dimensions:"
      ],
      "id": "jdPvixK_xsm5"
    },
    {
      "cell_type": "code",
      "execution_count": 23,
      "metadata": {
        "id": "6uaa60xcxr0h"
      },
      "outputs": [],
      "source": [
        "B = torch.rand(3, 4, 5)\n",
        "B_T = B[i, j].order(j, i)\n",
        "assert torch.allclose(B.permute(1, 0, 2), B_T)"
      ],
      "id": "6uaa60xcxr0h"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4d84ec1a"
      },
      "source": [
        "[^4] `order` is actually just a synonym for the already-existing `permute` method, which takes a list a dimension specifiers and puts the tensor in that order because rule #2 says that first-class dims can be passed as arguments to functions that previousely took only integers as dimensions. However, the name `permute` is confusing in this context since it implies dim objects have an original order, so we prefer to use `order` when writing code.\n",
        "\n",
        "Flattening and Splitting Dims\n",
        "-----------------------------\n",
        "\n",
        "**Tuples of dimensions** can be passed to both indexing and `order`. In indexing, this will split the dimension being indexed across the dimensions in the tuple.  In `order` it will flatten the dimensions in a single positional dimension:"
      ],
      "id": "4d84ec1a"
    },
    {
      "cell_type": "code",
      "execution_count": 24,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "abd5343e",
        "outputId": "ae8df95a-0ac1-45f8-b90b-a6fad4e149ab"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "3 2 4\n"
          ]
        }
      ],
      "source": [
        "i, j, k = dims(3)\n",
        "j.size = 2\n",
        "A = torch.rand(6, 4)\n",
        "a = A[(i, j), k] # split dim 0 into i,j\n",
        "print(i.size, j.size, k.size)"
      ],
      "id": "abd5343e"
    },
    {
      "cell_type": "code",
      "execution_count": 25,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "038801b7",
        "outputId": "4cc1f894-9b65-4558-d734-50112f1ae81f"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([3, 8])\n"
          ]
        }
      ],
      "source": [
        "r = a.order(i, (j, k)) # flatten j and k\n",
        "print(r.shape)"
      ],
      "id": "038801b7"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7f984ea2"
      },
      "source": [
        "The size of one unsized dimension in a tuple such as `i` can be inferred if the other sizes are known.\n",
        "\n",
        "Examples\n",
        "========\n",
        "\n",
        "The usefulness of dimension objects is best seen through examples. Let's look at some different ways they can be used.\n",
        "\n",
        "Einsum-style Products\n",
        "---------------------\n",
        "Rather than having [einsum](https://pytorch.org/docs/stable/generated/torch.einsum.html) as a custom operator, it is possible to express matrix products directly as a composition of multiplies and summations. The implementation will pattern match any multiplication followed by a sum to the right matrix-multiply operator."
      ],
      "id": "7f984ea2"
    },
    {
      "cell_type": "code",
      "execution_count": 26,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9cfdceec",
        "outputId": "dbe3280c-39c4-45c1-8289-29d42f9f1698"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([3, 5])"
            ]
          },
          "metadata": {},
          "execution_count": 26
        }
      ],
      "source": [
        "def mm(A, B):\n",
        "    i, j, k = dims(3)\n",
        "    r = (A[i, k] * B[k, j]).sum(k)\n",
        "    return r.order(i, j)\n",
        "mm(torch.rand(3, 4), torch.rand(4, 5)).shape"
      ],
      "id": "9cfdceec"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9c399ecb"
      },
      "source": [
        "The implementation of named tensors delays the execution of multiply to see if a summation follows it as it does above. If so, it will turn this pattern into the correct _optimized matrix product_, similar to how the `einsum` function works.\n",
        "\n",
        "Since it is no longer necessary to manually match math to matrix functions, other tensor products are easier to express, like the Gram matrix used in style transfer:"
      ],
      "id": "9c399ecb"
    },
    {
      "cell_type": "code",
      "execution_count": 27,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cd102da7",
        "outputId": "ae8aedc6-556f-4bff-f3bf-3740e2359739"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[[0.3218, 0.2184],\n",
              "         [0.2184, 0.3217]]])"
            ]
          },
          "metadata": {},
          "execution_count": 27
        }
      ],
      "source": [
        "def gram_matrix_new(y):\n",
        "    b, c, c2, h, w = dims()\n",
        "    r = (y[b, c, h, w] * y[b, c2, h, w]).sum((h, w))\n",
        "    r = r / (h.size * w.size)\n",
        "    return r.order(b, c, c2)\n",
        "  \n",
        "gram_matrix_new(torch.rand(1, 2, 3, 4))\n",
        "# [example adapted from http://einops.rocks/pytorch-examples.html]"
      ],
      "id": "cd102da7"
    },
    {
      "cell_type": "markdown",
      "source": [
        "Attention is another example that has several matrix products embedded inside it:"
      ],
      "metadata": {
        "id": "hGy8aF0Rp289"
      },
      "id": "hGy8aF0Rp289"
    },
    {
      "cell_type": "code",
      "execution_count": 28,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "188ad546",
        "outputId": "9a6c5eab-63de-47f6-884e-b6ece882a640"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[[0.6820, 0.6711, 0.6799, 0.6696],\n",
              "         [0.4203, 0.4092, 0.4210, 0.4056],\n",
              "         [0.6346, 0.6375, 0.6498, 0.6443],\n",
              "         [0.0051, 0.0013, 0.6346, 0.4695],\n",
              "         [0.3885, 0.7928, 0.3323, 0.8644],\n",
              "         [0.1283, 0.0595, 0.2990, 0.2795]],\n",
              "\n",
              "        [[0.5786, 0.5764, 0.5727, 0.5741],\n",
              "         [0.6020, 0.5996, 0.5899, 0.6014],\n",
              "         [0.3854, 0.3995, 0.4026, 0.4072],\n",
              "         [0.2443, 0.5288, 0.5756, 0.6316],\n",
              "         [0.1350, 0.4300, 0.6357, 0.5442],\n",
              "         [0.2228, 0.4503, 0.9129, 0.4328]]])"
            ]
          },
          "metadata": {},
          "execution_count": 28
        }
      ],
      "source": [
        "from torchdim import softmax\n",
        "def attention(K, Q, V):\n",
        "    batch, channel, key, query = dims(4)\n",
        "    k = K[batch, channel, key]\n",
        "    q = Q[batch, channel, query]\n",
        "    v = V[batch, channel, key]\n",
        "\n",
        "    a = (k * q).sum(channel) # matrix multiply\n",
        "    a = softmax(a * (channel.size ** -0.5), dim=key)\n",
        "    r = (v * a).sum(key) # matrix multiply\n",
        "    return torch.cat((r.order(batch, channel, query), Q), dim=1)\n",
        "\n",
        "inputs = (torch.rand(2, 3, 4) for _ in range(3))\n",
        "attention(*inputs)\n",
        "# [example adapted from http://einops.rocks/pytorch-examples.html]"
      ],
      "id": "188ad546"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c5135650"
      },
      "source": [
        "Reshaping tensors (einops) \n",
        "--------------------------\n",
        "\n",
        "Lots of operations in deep learning are just different ways of reshaping, splitting, and joining dimensions, such as the pixel shuffle used to upscale an image by turning channels into pixels:"
      ],
      "id": "c5135650"
    },
    {
      "cell_type": "code",
      "source": [
        "def pixel_shuffle(img, upscale_factor=2):\n",
        "    h2, w2, c, b, h, w = dims(6)\n",
        "    h2.size = w2.size = upscale_factor\n",
        "    return img[b, (c, h2, w2), h, w].order(b, c, (h, h2), (w, w2))\n"
      ],
      "metadata": {
        "id": "nx5bsjnmpDGC"
      },
      "id": "nx5bsjnmpDGC",
      "execution_count": 29,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f581561c"
      },
      "source": [
        "[Einops](http://einops.rocks) is an extension to einsum that adds support for the manipulation of dimensions through a few custom operators such as `rearrange`:"
      ],
      "id": "f581561c"
    },
    {
      "cell_type": "code",
      "execution_count": 30,
      "metadata": {
        "id": "ce513dcb"
      },
      "outputs": [],
      "source": [
        "def pixel_shuffle_einops(img, upscale_factor=2):\n",
        "    from einops import rearrange\n",
        "    return rearrange(img, 'b (c h2 w2) h w -> b c (h h2) (w w2)', h2=upscale_factor, w2=upscale_factor)"
      ],
      "id": "ce513dcb"
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "Named tensors with first-class dimensions can accomplish the same goal, but using PyTorch's existing operator set."
      ],
      "metadata": {
        "id": "Ew-dgu4nrnaK"
      },
      "id": "Ew-dgu4nrnaK"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "98047f30"
      },
      "source": [
        "Automatically batching Code (`vmap`, `xmap`)\n",
        "-----------------------------\n",
        "\n",
        "The implicit batching of Rule #1 means it is easy to created batched versions of existing PyTorch code. Simply bind a dim to the dimensions that should act as a batch, and then pass the tensor to the unbatched function. Since the unbatched function does not know about the dim, the dim will be implicictly batched over:"
      ],
      "id": "98047f30"
    },
    {
      "cell_type": "code",
      "execution_count": 31,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "eda7b7f2",
        "outputId": "4dd5d270-b75c-4714-e392-db0fc4c62bdd"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tensor([0.4775, 0.0000, 0.3423])\n",
            "with dims=(batch,) sizes=(3,)\n"
          ]
        }
      ],
      "source": [
        "batch_size, feature_size = 3, 5\n",
        "weights = torch.randn(feature_size)\n",
        "\n",
        "def model(feature_vec):\n",
        "    # Very simple linear model with activation\n",
        "    assert feature_vec.dim() == 1\n",
        "    return feature_vec.dot(weights).relu()\n",
        "\n",
        "examples = torch.randn(batch_size, feature_size)\n",
        "batch = dims(1)\n",
        "r = model(examples[batch])\n",
        "print(r)\n",
        "# in functorch: result = functorch.vmap(model)(examples)"
      ],
      "id": "eda7b7f2"
    },
    {
      "cell_type": "markdown",
      "source": [
        "This pattern also composes well with other code that also uses first class dimensions. For instance, we can write batched matrix multiply `bmm` by batching the `mm` operator.\n",
        "\n",
        "It doesn't matter whether the implementation of the function uses dimension objects, it is also possible to add additional batch dimensions and then call a function:"
      ],
      "metadata": {
        "id": "1fmWQoWSszB0"
      },
      "id": "1fmWQoWSszB0"
    },
    {
      "cell_type": "code",
      "source": [
        "def bmm(A, B):\n",
        "    i = dims(1) # note: i here is a different value from i inside mm so it works\n",
        "    return mm(A[i], B[i]).order(i)"
      ],
      "metadata": {
        "id": "4Ime0UoOs1TX"
      },
      "id": "4Ime0UoOs1TX",
      "execution_count": 32,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b91141cc"
      },
      "source": [
        "The equivalent code in JAX, using [xmap or vmap](https://jax.readthedocs.io/en/latest/notebooks/quickstart.html#auto-vectorization-with-vmap) are transforms over functions. So there is a lot of syntactic distance between the specification of the dimension mappings, and the values where those mappings apply. Dims express the mapping as indexing of the tensor, right at the place where the function is being applied.\n",
        "\n",
        "\n",
        "[xmap examples](https://jax.readthedocs.io/en/latest/notebooks/xmap_tutorial.html):\n",
        "```\n",
        "in_axes = [['inputs', 'hidden', ...],\n",
        "           ['hidden', 'classes', ...],\n",
        "           ['batch', 'inputs', ...],\n",
        "           ['batch', ...]]\n",
        "\n",
        "loss = xmap(named_loss, in_axes=in_axes, out_axes=[...])\n",
        "print(loss(w1, w2, images, labels))\n",
        "```\n",
        "\n",
        "Equivalent with dimension objects:\n",
        "\n",
        "```\n",
        "batch, inputs, hidden, classes = dims(4)\n",
        "print(loss(w1[inputs, hidden], w2[hidden, classes], images[batch, inputs], labels[batch]))\n",
        "```\n",
        "\n"
      ],
      "id": "b91141cc"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "afa0e95d"
      },
      "source": [
        "Composing matrix products, reshaping, and batching:\n",
        "---------------------\n",
        "\n",
        "Multi-headed attention is a good example of how these different uses compose. It reshapes the inputs, splitting out different attention heads. It batches over those attention heads, and it uses matrix products to compute attention scores."
      ],
      "id": "afa0e95d"
    },
    {
      "cell_type": "code",
      "execution_count": 33,
      "metadata": {
        "id": "acc1edbc"
      },
      "outputs": [],
      "source": [
        "from torchdim import softmax\n",
        "def multiheadattention(q, k, v, num_attention_heads, dropout_prob, use_positional_embedding):\n",
        "    batch, query_sequence, key_sequence, heads, features = dims(5)\n",
        "    heads.size = num_attention_heads\n",
        "\n",
        "    # binding dimensions, and unflattening the heads from the feature dimension\n",
        "    q = q[batch, query_sequence, [heads, features]]\n",
        "    k = k[batch, key_sequence, [heads, features]]\n",
        "    v = v[batch, key_sequence, [heads, features]]\n",
        "\n",
        "    # einsum-style operators to calculate scores,\n",
        "    attention_scores = (q*k).sum(features) * (features.size ** -0.5)\n",
        "\n",
        "    # use first-class dim to specify dimension for softmax\n",
        "    attention_probs = softmax(attention_scores, dim=key_sequence)\n",
        "\n",
        "    # dropout work pointwise, following Rule #1\n",
        "    attention_probs = torch.nn.functional.dropout(attention_probs, p=dropout_prob)\n",
        "\n",
        "    # another matrix product\n",
        "    context_layer = (attention_probs*v).sum(key_sequence)\n",
        "\n",
        "    # flatten heads back into features\n",
        "    return context_layer.order(batch, query_sequence, [heads, features])"
      ],
      "id": "acc1edbc"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5ec130b9"
      },
      "source": [
        "Indexing\n",
        "--------\n",
        "\n",
        "Rule #3 enables indexing because dimensions act as loop indices when used as a tensor. This allows for a lot of powerful behavior. The simplest might be using the dimensions to compute masks, such as extracing the upper triangular part of a matrix:"
      ],
      "id": "5ec130b9"
    },
    {
      "cell_type": "code",
      "execution_count": 34,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5e3bd826",
        "outputId": "459d8f63-59af-42f1-efb0-8c011ba4cdb5"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[0.0621, 0.0781, 0.6571, 0.4196],\n",
              "        [0.0000, 0.2611, 0.5586, 0.1965],\n",
              "        [0.0000, 0.0000, 0.7161, 0.8301]])"
            ]
          },
          "metadata": {},
          "execution_count": 34
        }
      ],
      "source": [
        "from torch import where\n",
        "def triu(A):\n",
        "   i,j = dims()\n",
        "   a = A[i, j]\n",
        "   return where(i <= j, a, 0).order(i, j)\n",
        "triu(torch.rand(3, 4))"
      ],
      "id": "5e3bd826"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "67478c7d"
      },
      "source": [
        "Embedding bag does an embedding table lookup followed by a sum, which can be expressed concisely:"
      ],
      "id": "67478c7d"
    },
    {
      "cell_type": "code",
      "execution_count": 35,
      "metadata": {
        "id": "cd02283d",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "21faf3ed-1285-4b25-ac44-9632a04ecd43"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[1.3716, 2.7679]])"
            ]
          },
          "metadata": {},
          "execution_count": 35
        }
      ],
      "source": [
        "def embedding_bag(input, embedding_weights):\n",
        "    batch, sequence, features = dims(3)\n",
        "    r = embedding_weights[input[batch, sequence], features].sum(sequence)\n",
        "    return r.order(batch, features)\n",
        "\n",
        "input = torch.tensor([[1, 0, 4, 3]])\n",
        "W = torch.rand(5,2)\n",
        "embedding_bag(input, W)"
      ],
      "id": "cd02283d"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "257a080b"
      },
      "source": [
        "Relative positional embeddings associate an embedding vector with the distance between the query and the key in the sequence.\n",
        "For instance, a key 3 and query 5 will have embedding ID `(5-3)=2`. We can use first-class dimensions to do the indexing arithmetic, and the embedding lookup:"
      ],
      "id": "257a080b"
    },
    {
      "cell_type": "code",
      "execution_count": 36,
      "metadata": {
        "id": "373f43b9"
      },
      "outputs": [],
      "source": [
        "def relative_positional_embedding(q, k, distance_embedding_weight):\n",
        "    batch, query_sequence, key_sequence, heads, features = dims(5)\n",
        "    q = q[batch, query_sequence, [heads, features]]\n",
        "    k = k[batch, key_sequence, [heads, features]]\n",
        "\n",
        "    distance = query_sequence - key_sequence\n",
        "    n_embeddings = distance_embedding_weight.size(0)\n",
        "    index_bias = n_embeddings // 2\n",
        "\n",
        "    assert key_sequence.size + bias <= n_embeddings\n",
        "\n",
        "    # indexing with dims\n",
        "    positional_embedding = distance_embedding_weight[distance + index_bias, features]\n",
        "\n",
        "    # matrix multiplies with dims\n",
        "    relative_position_scores_query = (q*positional_embedding).sum(features)\n",
        "    relative_position_scores_key = (k*positional_embedding).sum(features)\n",
        "    return  (relative_position_scores_query + relative_position_scores_key).order(batch, heads, key_sequence, query_sequence)"
      ],
      "id": "373f43b9"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4b29adb2"
      },
      "source": [
        "Tensor Puzzlers\n",
        "===============\n",
        "\n",
        "[Tensor Puzzlers](https://github.com/srush/Tensor-Puzzles), created by Sasha Rush, are a good exercise for learning the numpy and torch APIs by figuring out how to define common operations using a small set of primitive tensor operations.\n",
        "\n",
        "However, the difficulty of many of the puzzlers lies not in how to compute the answer but the awkwardness of the primitives themselves.\n",
        "\n",
        "**With first class dimensions, these puzzlers are nearly the same as the spec that defines them**\n",
        "\n",
        "\n",
        "### Puzzle 3 - outer\n",
        "\n",
        "Compute [outer](https://numpy.org/doc/stable/reference/generated/numpy.outer.html) - the outer product of two vectors."
      ],
      "id": "4b29adb2"
    },
    {
      "cell_type": "code",
      "execution_count": 37,
      "metadata": {
        "id": "83b1576b"
      },
      "outputs": [],
      "source": [
        "def outer_spec(a, b, out):\n",
        "    for i in range(len(out)):\n",
        "        for j in range(len(out[0])):\n",
        "            out[i][j] = a[i] * b[j]\n",
        "\n",
        "def outer(a, b):\n",
        "    i, j = dims(2)\n",
        "    return (a[i] * b[j]).order(i, j)"
      ],
      "id": "83b1576b"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ab1ce443"
      },
      "source": [
        "### Puzzle 4 - diag\n",
        "\n",
        "Compute [diag](https://numpy.org/doc/stable/reference/generated/numpy.diag.html) - the diagonal vector of a square matrix."
      ],
      "id": "ab1ce443"
    },
    {
      "cell_type": "code",
      "execution_count": 38,
      "metadata": {
        "id": "a69ab8fe"
      },
      "outputs": [],
      "source": [
        "def diag_spec(a, out):\n",
        "    for i in range(len(a)):\n",
        "        out[i] = a[i][i]\n",
        "\n",
        "def diag(a):\n",
        "    i = dims(1)\n",
        "    return a[i, i].order(i)"
      ],
      "id": "a69ab8fe"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ca63c7f9"
      },
      "source": [
        "### Puzzle 5 - eye\n",
        "\n",
        "Compute [eye](https://numpy.org/doc/stable/reference/generated/numpy.eye.html) - the identity matrix."
      ],
      "id": "ca63c7f9"
    },
    {
      "cell_type": "code",
      "execution_count": 39,
      "metadata": {
        "id": "0bf3b440"
      },
      "outputs": [],
      "source": [
        "from torch import where\n",
        "def eye_spec(out):\n",
        "    for i in range(len(out)):\n",
        "        out[i][i] = 1\n",
        "\n",
        "def eye(j: int):\n",
        "    i,j = dims(sizes=[j, j])\n",
        "    return where(i == j, 1, 0).order(i, j)"
      ],
      "id": "0bf3b440"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "223a9b57"
      },
      "source": [
        "### Puzzle 6 - triu\n",
        "\n",
        "Compute [triu](https://numpy.org/doc/stable/reference/generated/numpy.triu.html) - the upper triangular matrix."
      ],
      "id": "223a9b57"
    },
    {
      "cell_type": "code",
      "execution_count": 40,
      "metadata": {
        "id": "9ec0cfc0"
      },
      "outputs": [],
      "source": [
        "def triu_spec(out):\n",
        "    for i in range(len(out)):\n",
        "        for j in range(len(out)):\n",
        "            if i <= j:\n",
        "                out[i][j] = 1\n",
        "            else:\n",
        "                out[i][j] = 0\n",
        "\n",
        "def triu(j: int):\n",
        "    i,j = dims(sizes=[j, j])\n",
        "    return where(i <= j, 1, 0).order(i, j)"
      ],
      "id": "9ec0cfc0"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "17bf4f6d"
      },
      "source": [
        "### Puzzle 8 - diff\n",
        "\n",
        "Compute [diff](https://numpy.org/doc/stable/reference/generated/numpy.diff.html) - the running difference."
      ],
      "id": "17bf4f6d"
    },
    {
      "cell_type": "code",
      "execution_count": 41,
      "metadata": {
        "id": "1df20624"
      },
      "outputs": [],
      "source": [
        "def diff_spec(a, out):\n",
        "    out[0] = a[0]\n",
        "    for i in range(1, len(out)):\n",
        "        out[i] = a[i] - a[i - 1]\n",
        "def diff(a, i: int):\n",
        "    i = dims(1)\n",
        "    d = a[i] - a[i - 1]\n",
        "    return where(i - 1 >= 0, d, a[i]).order(i)"
      ],
      "id": "1df20624"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e92ddcb8"
      },
      "source": [
        "### Puzzle 9 - vstack\n",
        "\n",
        "Compute [vstack](https://numpy.org/doc/stable/reference/generated/numpy.vstack.html) - the matrix of two vectors"
      ],
      "id": "e92ddcb8"
    },
    {
      "cell_type": "code",
      "execution_count": 42,
      "metadata": {
        "id": "4855791c"
      },
      "outputs": [],
      "source": [
        "def vstack_spec(a, b, out):\n",
        "    for i in range(len(out[0])):\n",
        "        out[0][i] = a[i]\n",
        "        out[1][i] = b[i]\n",
        "\n",
        "def vstack(a, b):\n",
        "    v, i = dims(sizes=[2, None])\n",
        "    return where(v == 0,  a[i], b[i]).order(v, i)"
      ],
      "id": "4855791c"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b66f5043"
      },
      "source": [
        "### Puzzle 10 - roll\n",
        "\n",
        "Compute [roll](https://numpy.org/doc/stable/reference/generated/numpy.roll.html) - the vector shifted 1 circular position."
      ],
      "id": "b66f5043"
    },
    {
      "cell_type": "code",
      "execution_count": 43,
      "metadata": {
        "id": "4be581a0"
      },
      "outputs": [],
      "source": [
        "def roll_spec(a, out):\n",
        "    for i in range(len(out)):\n",
        "        if i + 1 < len(out):\n",
        "            out[i] = a[i + 1]\n",
        "        else:\n",
        "            out[i] = a[i + 1 - len(out)]\n",
        "\n",
        "def roll(a, i: int):\n",
        "    i = dims(sizes=[a.size(0)])\n",
        "    return a[where(i + 1 < i.size, i + 1, 0)].order(i)"
      ],
      "id": "4be581a0"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "935e569e"
      },
      "source": [
        "### Puzzle 11 - flip\n",
        "\n",
        "Compute [flip](https://numpy.org/doc/stable/reference/generated/numpy.flip.html) - the reversed vector"
      ],
      "id": "935e569e"
    },
    {
      "cell_type": "code",
      "execution_count": 44,
      "metadata": {
        "id": "16ec4c8b"
      },
      "outputs": [],
      "source": [
        "def flip_spec(a, out):\n",
        "    for i in range(len(out)):\n",
        "        out[i] = a[len(out) - i - 1]\n",
        "\n",
        "def flip(a, i: int):\n",
        "    i = dims(sizes=[a.size(0)])\n",
        "    return a[i.size - i - 1].order(i)"
      ],
      "id": "16ec4c8b"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ac33ea48"
      },
      "source": [
        "### Puzzle 14 - sequence_mask\n",
        "\n",
        "\n",
        "Compute [sequence_mask](https://www.tensorflow.org/api_docs/python/tf/sequence_mask) - pad out to length per batch."
      ],
      "id": "ac33ea48"
    },
    {
      "cell_type": "code",
      "execution_count": 45,
      "metadata": {
        "id": "566db352"
      },
      "outputs": [],
      "source": [
        "def sequence_mask_spec(values, length, out):\n",
        "    for i in range(len(out)):\n",
        "        for j in range(len(out[0])):\n",
        "            if j < length[i]:\n",
        "                out[i][j] = values[i][j]\n",
        "            else:\n",
        "                out[i][j] = 0\n",
        "\n",
        "def sequence_mask(values, length):\n",
        "    j, i = dims()\n",
        "    v = values[i, j]\n",
        "    return where(j < length[i], v, 0).order(i, j)"
      ],
      "id": "566db352"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4dc80f7a"
      },
      "source": [
        "Advantages of First-class Dimensions over String Dimensions\n",
        "===================================================================\n",
        "\n",
        "The most prominent difference between named tensors using first-class dimensions and alternatives (einops, named tensors implemented in PyTorch today , [tensors considered harmful](https://nlp.seas.harvard.edu/NamedTensor), or xmap) is that dimensions are objects rather than strings. Using objects has a number of nice properties.\n",
        "\n",
        "### Avoiding naming conflicts\n",
        "\n",
        "Using strings for dimensions introduces the possibility that two unrelated dimensions are given the same name. Using objects instead makes it clear the same names are not the same dimension. It's like the difference between having only global variables, and having the ability to locally bind names in functions.\n",
        " For instance, we defined `bmm` by batching a call to `mm`, and even though they both use the name `i` to identify a dimension.  Because each `i` is a different object, there is no naming conflict:"
      ],
      "id": "4dc80f7a"
    },
    {
      "cell_type": "code",
      "execution_count": 46,
      "metadata": {
        "id": "b2859d73"
      },
      "outputs": [],
      "source": [
        "def mm(A, B):\n",
        "    i, j, k = dims()\n",
        "    r = (A[i, k] * B[k, j]).sum(k)\n",
        "    return r.order(i, j)\n",
        "\n",
        "def bmm(A, B):\n",
        "    i = dims() # note: doesn't matter than mm internally also uses i\n",
        "    return mm(A[i], B[i])"
      ],
      "id": "b2859d73"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3fa18092"
      },
      "source": [
        "Einops avoids conflicts by ensuring names are all introduced and removed in a single expression, but this precludes using long-lived dimensions to present implicit batching similar to xmap. When nested, JAX's xmap seems to consider axes the same if the string name matches. In the above example it would consider the `i` dimension to be the same dimension in both `bmm` and `mm` so the code would error.\n",
        "\n",
        "\n",
        "### Reuse the same operator set\n",
        "\n",
        "Having a new object type allows us to extend the existing operator set of PyTorch rather than come up with new operators. For instance, binding dimensions using indexing follows semantically from Rules #1 and #3, so there is no need for a special operator to do binding. Even unbinding is just the `permute` operator which follows from Rule #2, though we call it `order` for clarity. In contrast, using strings requires coming up with new APIs such as `einsum` for matrix multiplies, or `rearrange` for doing permutations.\n",
        "\n",
        "### Allows dims to act as tensors\n",
        "\n",
        "Rule #3 is not possible with strings since we cannot make strings behave as tensors. Without this rule, all of the indirect indexing that dims enable would not be easy to express.\n",
        "\n",
        "### Dims can have methods\n",
        "For instance, as objects, dims can have a size, which allows us to do size inference of dimensions in various places in the API where string based APIs would have to take additional arguments specifying size.\n",
        "\n",
        "\n",
        "Comparison to tensor compilers or languages (e.g. TVM or Dex)\n",
        "=============================================================\n",
        "\n",
        "The semantics and surface syntax of dimension objects resembles the kind of code written in tensor compilers such as [Halide](https://halide-lang.org), [TVM](https://tvm.apache.org), [Tensor Comprehensions](https://github.com/facebookresearch/TensorComprehensions), or the language [Dex](https://github.com/google-research/dex-lang). \n",
        "\n",
        "These compilers and language have syntax and semantics that resemble the loop-level analogy similar to first-class dimensions. However, as compilers or statically typed languages, they require some binding code to go from running deep learning framework code in Python to using the compiled language. This often at least requires refactoring the compiled parts into their own functions, and may require defining a gradient function. Similar to graph mode frameworks, this adds friction to using and debugging the code.\n",
        "\n",
        "Dimension objects are just an extension of the existing PyTorch tensors and eager sematics, so there is no friction switching between normal Python code and code that uses them. However, since loops over the dimensions are defined implicitly, they can still execute in Python with good performance compared to explicit loops. Furthermore, with dimension objects, a tensors containing dimensions can compute through code that is oblivous to the dimension such as batching examples. There is no need to separate code into 'compiled' vs 'eager'.\n",
        "\n",
        "In this way, first-class dims are a way of adapting the nicer syntax of these array compilers and languages to eager numpy-style libraries.\n",
        "\n",
        "\n",
        "Performance Expectations\n",
        "========================\n",
        "First-class dimensions are not a compiler. They provide syntax for existing PyTorch operations such as advanced indexing that is easier to read and write. For large sized tensors, the performance of any statements including them will be the same as using the already existing operations. An important exception is the pattern matching of products and summation, where performance will be improved by issuing to a matrix-multiply kernel. The C++ implementation of dimensions adds a small overhead of around 2us on top of PyTorch's normal overhead of 8us to each function that uses them. In the future, the implementation can encorporate more fusion optimization to further improve performance of this style of code."
      ],
      "id": "3fa18092"
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "afa0e95d",
        "5ec130b9",
        "ab1ce443",
        "ca63c7f9",
        "223a9b57",
        "17bf4f6d",
        "e92ddcb8",
        "b66f5043",
        "935e569e",
        "ac33ea48",
        "3fa18092"
      ],
      "name": "torchdim.ipynb",
      "provenance": []
    },
    "jupytext": {
      "cell_metadata_filter": "-all",
      "formats": "md:myst",
      "text_representation": {
        "extension": ".md",
        "format_name": "myst",
        "format_version": 0.13,
        "jupytext_version": "1.13.8"
      }
    },
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}