{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$\\color{brown}{Preamble}$: At the time of this writing, I'm using **PyTorch** **`v1.7.1`** binded with **`cuda11.0`** and **`cudnn8.0`**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"version: \", torch.__version__)\n",
    "mydevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(\"device : \", mydevice)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Einstein Summation (einsum)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Einsum is a powerful concept for processing tensors while at the same time writing very succinct code. The reasons to adopt **_einsum_** are:\n",
    "\n",
    "  - the code is usually one-liner\n",
    "  - it's memory efficient\n",
    "  - less error-prone"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us now see some sample problems to fully grasp the power of einsum.\n",
    "\n",
    "#### **inputs**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# some input tensors to work with\n",
    "\n",
    "vec = torch.tensor([0, 1, 2, 3])\n",
    "\n",
    "aten = torch.tensor([[11, 12, 13, 14],\n",
    "                     [21, 22, 23, 24],\n",
    "                     [31, 32, 33, 34],\n",
    "                     [41, 42, 43, 44]])\n",
    "\n",
    "bten = torch.tensor([[1, 1, 1, 1],\n",
    "                     [2, 2, 2, 2],\n",
    "                     [3, 3, 3, 3],\n",
    "                     [4, 4, 4, 4]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "-------------\n",
    "### matrix multiplication\n",
    "\n",
    "  $\\textbf{C}_{ik}$  =  $\\sum_{j}$ $\\textbf{A}_{i\\color{green}{j}}$ * $\\textbf{B}_{\\color{green}{j}k}$\n",
    "\n",
    "For a matrix multiplication to work, the number of columns in the first matrix (e.g., $A$) should match the number of rows in the second matrix (e.g., $B$)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c = torch.einsum('ij, jk -> ik', aten, bten)\n",
    "print(\"einsum matmul: \\n\", c)\n",
    "\n",
    "# sanity check\n",
    "c = torch.matmul(aten, bten)   # or: aten.mm(bten)\n",
    "print(\"torch matmul: \\n\", c)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "------------------\n",
    "### hadamard product (*i.e.,* element-wise product of tensors)\n",
    "$\\textbf{C}_{ij}$  =  $\\textbf{A1}_{ij}$ * $\\textbf{A2}_{ij}$ * ... *  $\\textbf{AN}_{ij}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hp = torch.einsum('ij, ij, i -> ij', aten, bten, vec)     # note: `vec` is treated as a column vector\n",
    "print(\"einsum hadamard product: \\n\", hp)\n",
    "\n",
    "# sanity check\n",
    "ep = aten * bten * vec[:, None]\n",
    "print(\"element-wise product: \\n\", ep)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$\\color{brown}{Note}$: we can raise the elements of a tensor to power `n` by repeating the tensor `n` times. For instance, a tensor can be *cubed* by repeating it 3 times."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hp = torch.einsum('ij, ij, ij -> ij', bten, bten, bten)\n",
    "print(\"einsum hadamard product: \\n\", hp)\n",
    "\n",
    "# sanity check\n",
    "ep = bten * bten * bten\n",
    "print(\"element-wise product: \\n\", ep)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}