Taco Tutorial via MTTKRP.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction to TACO\n",
    "\n",
    "TACO, which stands for Tensor Algebra Compiler, is a library for performing sparse and dense linear algebra and tensor algebra computations. \n",
    "\n",
    "The computations can range from relatively simple ones like sparse matrix-vector multiplication to more complex ones like matricized tensor times Khatri-Rao product. All these computations can be performed on any mix of dense and sparse tensors. Under the hood, TACO automatically generates efficient code to perform these computations.\n",
    "\n",
    "This notebook provides a brief introduction to the TACO Python library. For a more comprehensive overview, please see the documentation linked [here](http://tensor-compiler.org/symphony-docs/index.html). We will also link to relevant pages as we progress."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Table of Contents:\n",
    "* [Getting Started](#first-section)\n",
    "* [Defining Tensor Formats](#second-section)\n",
    "* [NumPy and SciPy I/O](#third-section)\n",
    "* [Example Application: MTTKRP](#fourth-section)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Getting Started <a class=\"anchor\" id=\"first-section\"></a>\n",
    "\n",
    "First, let's import TACO. Press `Shift` + `Enter` to run the code below. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pytaco as pt\n",
    "from pytaco import dense, compressed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the above, `dense` and `compressed` are [mode (dimension) formats](http://tensor-compiler.org/symphony-docs/reference/rst_files/mode_format.html). We can think of tensors as multi-dimensional arrays, and the mode formats allow us to specify how we would like to store the data in each dimension: \n",
    "* If a dimension is `dense`, then all of the elements in that dimension are stored. \n",
    "* And if a dimension is `compressed`, then only nonzeros are stored.\n",
    "\n",
    "For example, we can declare a $512 \\times 64 \\times 2048$ [tensor](http://tensor-compiler.org/symphony-docs/reference/rst_files/tensor_class.html) whose first dimension is dense and second and third dimensions are compressed:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "T = pt.tensor([512, 64, 2048], pt.format([dense, compressed, compressed]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can initialize $T$ by calling its `insert` [method](http://tensor-compiler.org/symphony-docs/reference/rst_files/functions/pytaco.tensor.insert.html) to add a nonzero element to the tensor. The `insert` method takes two arguments: a list of coordinates and the value to be inserted at those coordinates:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set T(0, 1, 0) = 42.0\n",
    "T.insert([0, 1, 0], 42.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If multiple elements are inserted at the same coordinates, they are summed together:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set T(0, 0, 1) = 12.0 + 24.0 = 36.0\n",
    "T.insert([0, 0, 1], 12.0)\n",
    "T.insert([0, 0, 1], 24.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can then iterate over the nonzero elements of the tensor as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for coordinates, value in T: \n",
    "    print(\"Coordinates: {}, Value: {}\".format(coordinates, value))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Defining Tensor Formats <a class=\"anchor\" id=\"second-section\"></a>\n",
    "\n",
    "Consider a matrix $M$ (aka a two-dimensional tensor) containing the following values:\n",
    "\n",
    "$$M = \\begin{bmatrix}\n",
    "    6 & \\cdot & 9 & 8 \\\\\n",
    "    \\cdot & \\cdot & \\cdot & \\cdot \\\\\n",
    "    5 & \\cdot & \\cdot & 7  \n",
    "\\end{bmatrix}$$\n",
    "\n",
    "Denote the rows and columns as dimensions $d_1$ and $d_2$, respectively. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We look at how $M$ is represented differently in different formats. For convenience, let's define a helper function to initialize $M$. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_example_matrix(format):\n",
    "    M = pt.tensor([3, 4], format)\n",
    "    M.insert([0, 0], 6)\n",
    "    M.insert([0, 2], 9)\n",
    "    M.insert([0, 3], 8)\n",
    "    M.insert([2, 0], 5)\n",
    "    M.insert([2, 3], 7)\n",
    "    return M"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## (dense $d_1$, dense $d_2$)\n",
    "Note that passing in `dense` makes all of the dimensions dense. This is equivalent to `pt.format([dense, dense])`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "make_example_matrix(dense)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For this example, we focus on the last line of the output, the `vals` array: since all values are stored, it is a flattened $3 \\times 4$ matrix stored in row-major order."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## (dense $d_1$, compressed $d_2$)\n",
    "This is called compressed sparse row (CSR) format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "csr = pt.format([dense, compressed])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "make_example_matrix(csr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since $d_1$ is dense, we need only store the `size` of the dimension; values, both zero and nonzero, are stored for every coordinate in dimension $d_1$.  \n",
    "\n",
    "Since $d_2$ is compressed, we store a `pos` array (`[0, 3, 3, 5]`) and an `idx` array (`[0, 2, 3, 0, 3]`); these together form a segmented vector with one segment per entry in the previous dimension. The `idx` array stores all the indices with nonzero values in the dimension, while the `pos` array stores the location in the `idx` array where each segment begins. In particular, segment $i$ is stored in locations `pos[i]:pos[i+1]` in the `idx` array.\n",
    "\n",
    "The below animation visualizes the format. Hover over any non-empty entry of the matrix on the left to see how the value is stored."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run ./animation/animation_2.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## (dense $d_2$, compressed $d_1$)\n",
    "\n",
    "We switch the order of the dimensions by passing in `[1, 0]` for the optional parameter `mode_ordering`. This results in a column-major (rather than row-major) format called compressed sparse column (CSC)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "csc = pt.format([dense, compressed], [1, 0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "make_example_matrix(csc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this format, $d_2$ has only `size`, while $d_1$ has a `pos` and `idx` array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%run ./animation/animation_3.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## (compressed $d_1$, compressed $d_2$, compressed $d_3$)\n",
    "\n",
    "To more clearly visualize the compressed sparse fiber (CSF) format, where all dimensions are sparse, we move to three dimensions. The tensor $B$ defined below is a $2 \\times 3 \\times 4$ tensor.\n",
    "\n",
    "Similarly as above, passing in `compressed` is equivalent to `pt.format([compressed, compressed, compressed])`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "B = pt.tensor([2, 3, 4], compressed)\n",
    "\n",
    "# First layer.\n",
    "B.insert([0, 0, 0], 6)\n",
    "B.insert([0, 0, 2], 9)\n",
    "B.insert([0, 0, 3], 8)\n",
    "B.insert([0, 2, 0], 5)\n",
    "B.insert([0, 2, 3], 7)\n",
    "\n",
    "# Second layer. \n",
    "B.insert([1, 0, 1], 2)\n",
    "B.insert([1, 1, 0], 3)\n",
    "B.insert([1, 2, 2], 6)\n",
    "B.insert([1, 1, 3], 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "B"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This animation represents $N$ as two $3 \\times 4$ matrices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run ./animation/animation_4.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# NumPy and SciPy I/O <a class=\"anchor\" id=\"third-section\"></a>\n",
    "\n",
    "We can also initialize tensors with NumPy arrays or SciPy sparse (CSR or CSC) matrices. Let's start by importing the packages we need."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy as sp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Given a NumPy array such as the one randomly generated below, we can convert it to a TACO tensor using the `pytaco.from_array` [function](http://tensor-compiler.org/symphony-docs/reference/rst_files/functions/pytaco.from_array.html), which creates a dense tensor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "array = np.random.uniform(size=10)\n",
    "tensor = pt.from_array(array)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "tensor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also export TACO tensors to NumPy arrays."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tensor.to_array()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Similarly, given a SciPy sparse matrix, we can convert it to a TACO tensor. For a CSR matrix like the one below, we use the `pytaco.from_sp_csr` [function](http://tensor-compiler.org/symphony-docs/reference/rst_files/functions/pytaco.from_sp_csr.html), which creates a CSR tensor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "size = 100\n",
    "density = 0.1\n",
    "sparse_matrix = sp.sparse.rand(size, size, density, format = 'csr')\n",
    "tensor = pt.from_sp_csr(sparse_matrix)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And we can export the tensor to a SciPy sparse matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tensor.to_sp_csr()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example Application: MTTKRP <a class=\"anchor\" id=\"fourth-section\"></a>\n",
    "\n",
    "The following example demonstrates how computations on tensors are performed using TACO."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Matricized tensor times Khatri-Rao product (MTTKRP) is a bottleneck operation in various algorithms — such as Alternating Least Squares — for computing sparse tensor factorizations like the Canonical Polyadic Decomposition. Mathematically, mode-1 MTTKRP (for three-dimensional tensors) can be expressed as\n",
    "\n",
    " $$A = B_{(1)}(D\\odot C),$$ \n",
    "\n",
    "where $A$, $C$, and $D$ are typically dense matrices, $B$ is a three-dimensional tensor (matricizied along the first mode), and $\\odot$ denotes the Khatri-Rao product. This operation can also be expressed in [index notation](http://tensor-compiler.org/symphony-docs/pycomputations/index.html#specifying-tensor-algebra-computations) as \n",
    "\n",
    "$$A_{ij} = B_{ikl} \\cdot D_{lj} \\cdot C_{kj}.$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Starting with the $B$ generated above in CSF format, we can view its `shape` attribute, which we expect to be `[2, 3, 4]`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "B.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Examining the formula, we choose compatible dimensions for $C$ and $D$ and generate them randomly with NumPy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "otherSize = 6\n",
    "C = pt.from_array(np.random.uniform(size=(B.shape[1], otherSize)))\n",
    "D = pt.from_array(np.random.uniform(size=(B.shape[2], otherSize)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Expressing the Computation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can express the result $A$ as a dense matrix. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A = pt.tensor([B.shape[0], otherSize], dense)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The syntax for TACO computations closely mirrors index notation, with the caveat that we also have to explicitly declare the index variables beforehand:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "i, j, k, l = pt.get_index_vars(4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A[i, j] = B[i, k, l] * D[l, j] * C[k, j]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Performing the Computation\n",
    "\n",
    "Once a tensor algebra computation has been defined, we can simply invoke the result tensor's `evaluate` method to perform the actual computation.\n",
    "\n",
    "[Under the hood](http://tensor-compiler.org/symphony-docs/pycomputations/index.html#performing-the-computation), TACO will first invoke the result tensor's `compile` method to generate code that performs the computation. TACO will then perform the actual computation by first invoking `assemble` to compute the sparsity structure of the result and subsequently invoking `compute` to compute the values of the result's nonzero elements."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A.evaluate()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we define a computation and then access the result without first manually invoking `evaluate` or `compile`/`assemble`/`compute`, TACO will automatically invoke the computation immediately before the result is accessed."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can display the result. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}