{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Einops tutorial, part 1: basics\n",
    "\n",
    "<!-- <img src='http://arogozhnikov.github.io/images/einops/einops_logo_350x350.png' height=\"80\" /> -->\n",
    "\n",
    "## Welcome to einops-land!\n",
    "\n",
    "We don't write \n",
    "```python\n",
    "y = x.transpose(0, 2, 3, 1)\n",
    "```\n",
    "We write comprehensible code\n",
    "```python\n",
    "y = rearrange(x, 'b c h w -> b h w c')\n",
    "```\n",
    "\n",
    "\n",
    "`einops` supports widely used tensor packages (such as `numpy`, `pytorch`, `jax`, `tensorflow`), and extends them.\n",
    "\n",
    "## What's in this tutorial?\n",
    "\n",
    "- fundamentals: reordering, composition and decomposition of axes\n",
    "- operations: `rearrange`, `reduce`, `repeat`\n",
    "- how much you can do with a single operation!\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preparations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Examples are given for numpy. This code also setups ipython/jupyter\n",
    "# so that numpy arrays in the output are displayed as images\n",
    "import numpy\n",
    "from utils import display_np_arrays_as_images\n",
    "\n",
    "display_np_arrays_as_images()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load a batch of images to play with"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(6, 96, 96, 3) float64\n"
     ]
    }
   ],
   "source": [
    "ims = numpy.load(\"./resources/test_images.npy\", allow_pickle=False)\n",
    "# There are 6 images of shape 96x96 with 3 color channels packed into tensor\n",
    "print(ims.shape, ims.dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAGAAAABgCAIAAABt+uBvAAAG+klEQVR4Ae2ba0xUVxCAy2tRXiIoQlCKSqVBrVQXIoQgLYqCgmmsGxPTrDZUKmho3FgbGtH+qKZJSSMVf2BEjWhJZEVTZKEgj2KpRbCFgqDotvWBoBRY3r6g02xCVu7ee87eO3cVcwwxd2fmzJnzMfdwHrM2Y+3tb7B//ARs+VVM8z8BBoiQBwwQA0QgQFCzDGKACAQIapZBDBCBAEHNMogBIhAgqFkGMUAEAgQ1yyAGiECAoGYZxAARCBDULIMYIAIBgpplEANEIEBQswxigAgECGqWQQRA9gS9nOrHT578du3az1euNLa0tOn17Z2dA4ODIJw6ZYrT1KkzPDzm+vnN8/MLCQ4OVyrn+/vLGQuvbxvrXxyOjY2VVFbmarUXSkqACG9oLyremjtXlZCgVqng4UWNvJ+sCmh0dPRUfv43WVktbW3ihmVra7th7dp9u3YtDAwU58HSVtYDBG9TSlpafWOjpSFy7R0cHDRJSfs0mimOjlwtrsQagCBxDmRm7s/IeP78OWL0yiVLCnJyZvv4IPrkupId0NDw8MZt24ouXeL2LV3i7eVVqdUGzp8v3RWfB3n/zPcYDCtVKpnowJA6Hj6M3rjxrzt3+IYnXS4joJHHjxPU6l/r66VHKeDhfkfH+q1bh0dGBGykqGQE9NHOnZdra6UER9n2z5aWz9LTKY0tNZMLUNbx4/mFhZZGI9o+Ozf3l6tXRTcXaCjLJN1669bSmBj50t7seN5dtKi+pMTGxsasVrRQlgxK3bvXynRg/L83NV0sKxMNgq8hfgYVlpbGq9V8/XHlbq6uG+LiVkZGLl282HP6dHc3N0N/f1d3d1NrK/z5K9Dpevv6uK3MSiKXL686d86sSrQQH5ByzRrK5bKdnd3u7ds/T0mZPm0a3wBgofD1oUPfZWfDapPPxlTeVlMTgLqtRX7Fyi9fpqTj6uJy8dSpg2lpAnRg5KD9Nj09/+hRyl1F3vnzprykPyMD+j4nhyYmmEpzDx9eHRVFYww2H8TGZh08SGMMrySNGb0NJiBDX5+uvJym7082b06IiaGxHLf5eNOmmBUrxj/yPfzR3Axh8GlFyDEBnS8uhuMuYhCOCsV+jYZoxjXYnZzMFU6QwFSFuzrFBPRTVdWEcM1+jIuO9pk1y6xKWBgVFubi7CxsA1rKSZDox2iACaimro6mV5hQaMy4Nvb29rAU4MonSK7fvDlBIuUjGqDOR4/+vnuXJhQ4x6ExM2sza+ZMs3JToejjSlMn489oh/awvRh3KvwQRDHXCnsQ1sL+XtjAIi1aBlGmj0XBiTPu7u19+vSpuLbcVq8hILg1edTdzR2qOAkaINgTiItAjlaDQ0NYbtEAwdkzVkzS/YzgHTCiAbL++YYAR5r1qkBzUxUaIFOnr9MzGiC4TX91uCgUCqxg0ABBxQFWTNL9KBwcpDsxekADBMUYWDFJ9+Ps5CTdCTKgN2fPxopJuh8Pd3fpTowe0LYa/nPmUMbU0dBAs6Wi9Ca3GdorRl+PIutNMTovNECQ1QvmzaOJD6qnaMxeERs0QDCe8JAQmlEdOXmyf2CAxvJVsMEEFL9qFc2QHnZ1fbpnD2wpaYxfug0mIDhLhcscmiGdKShI1GgQDyVoOhVngwkIrq5U8fGUceTk5YXExkKJK6W9WTMoWYODcHVq6lcZGWYNpAuRb1bhuHNhVJRFr09EaOiupCS4I6PfrMDFTll19Y+lpXDN/W9PD1CAS6QLJ05Ix8H1gAwIOtiQmHiuqIjbk7AEsu/9iIjQ4OCgBQsCAwLgkt7FyQnuMOCQAO7mew0GAAG39XWNjXUNDfBrmFDuCFV4rdXVwl2I0+IDgmUOJJGVTz/gwmNYr4f/xVEQaIU5Bxm7gfL4L1NTBbqUQ/Xs2TO9PJWK+IBg/F/s2AHvixwgBHzeuH1bQCtaJQsgqGv54cgRP19f0WGJaHiD+t7JIueyAIIIvGbMKD971tfb26JopBhPpgwyjhO+n1Oh1eKWMwkQnHyAYDDwzZxanY6+Dkhg/ETVpAQEo4ISMd3p01kHDkAtInGQUgxgi0dfzUjfkVxzkGkEUE+WvGXL9aoqqJuCL+qYqnCf5Zin8ReKwmP+5969zGPH4Nt08AsXtqTXzvT0XL969Yfr1kVHRKCvFa0NyDhsWNcVV1QUlpXBVlPcASOUqYUple+Fh8OCK2zZMlhY0AO1yPLlADIN8W57OxQWNjQ339TroXLl/oMHcM0POxX4gU0v3E8YN2VwkAILq7cDAow/7wQFUda9mvYl4vnlAxIRtDWbWGOStuZ40PtigAhIGSAGiECAoGYZxAARCBDULIMYIAIBgpplEANEIEBQswxigAgECGqWQQwQgQBBzTKIASIQIKhZBjFABAIE9X+hFi2cRY+DcwAAAABJRU5ErkJggg==",
      "text/plain": []
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# display the first image (whole 4d tensor can't be rendered)\n",
    "ims[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAGAAAABgCAIAAABt+uBvAAAE8ElEQVR4Ae2cW0gUURjHW68Vrkn3vKSsVj6skS6ZlUahWBop5KXIaiGqh97KfPAlWfApCI2eIovA1Jd60oQiK+/hmqmIZUW0ubtBeemibqWrfbIkC+p8Z3TmnE0+Xxzm/58z5/vtf44z456jsU/Zl9HP/AS85pdImSbgIwTDu9fvKssqG+sarRbrL8evNevWxMTFHMo8lHUyy9fPV0iX5juphvMl9uf3n6KLReU3yycnJ2f3KWRzSOnd0r0H9s6WRO3hCgjCkpuS297SLlGtt7f3tdvXco25Eh6eEtcxKP9svjQdqNzpdBacKzA3m3lSkDgXvwR1mjvT49MluuIu7di5o7at1n2PqG1+Caq4VcFeJNDsedXD7lfPyQ9Q89NmWWU0PGmQ5VfJzA+Q7ZNNVg3Wj1ZZfpXMnACNjY6Nj4/LquHb8DdZfpXMnACtWLnCx0feTWngqkCVapbVLCdAGo1mU+gmWT0LDQ+V5VfJzAkQ9D4xOVFWDUkpSbL8Kpn5ATpx9gR7DfBott2wnd2vnpMfIEOCIeNYBkslMFqZSkxwVbKY1fbwAwSVlNwpidsVJ12Sl5fX1ZtXE/YlSNu4qVwBwd+y+8/u553LAwpzVggDedWjquNnjs+pCtnJ71nMvby3vW/hyaPpaZPNYnM4HNPvg2JjDmYezD6V7efv5+4Uvi0GkPCy2Tswd9TZj1/yTgKEfMQEiAAhBBCZEoQAkveELdHYZ+tnQ5hBwrAAqbq1Gu6/F3CggodQghCYBIgAIQQQmRJEgBACiEwJIkAIAUSmBBEghAAiU4IQQIo9asDbUsbvYg0NDOnX6ZF+eYxMCUI+CgJEgBACiEwJIkAIAUSmBBEghAAiU4IIEEIAkSlBBAghgMiUIAKEEEBkShABQgggMiWIACEEEFlAgjzk+70ImH+yAEAw6fLf2f+D3wIA+fiyvgh3TjiFIxQAyH+5P2PZoyOjjE71bAIAwSUWoA1gKWngywCLTVWPAEBQT9DqIJaq+nr6WGyqesQA2hC8gaWqluctLDZVPWIAMU6Wg7nP79+8V7V+tHExgHRbdGjPXAZTvonRqZJNDKDomGjGeupq6wovFKITgicmJuof1xecLzAeMTK2zGgTM5kFFn2Jj4hn7CLYIiIjjBeMSclJwWHB2lVaWAPk5/eftn4bTB3v7e7teNEBF6PrnkAbqO37ruTQLgYQ1Lw7crflg4WdEbuzw9qxMWQju1/aKeYSgz6lZqRK92zBKkxGW/Cxsw8UBijndM7s3iiyZ4kA0sfqVZqYCstbKQLa1YiwBMHpL5suK1jJTFNLJEFQz579e2CS6kxhSm0sHUBApPhGcVR0lFJoXO0MDw4Pfh1Uqk2RlxjUACuYlD8sh7sbpepxtaNgiAQDgnrCdeE1rTUKLkQBCw0q+CJJPCBgBPd1MHcOxmxYeWAxUYIkXrpyyWwxpxxOWUw77scKu5N278TMNrwhK7te9uDeA1nLVYVFhAGRtKNpMOrPt6bDzCnkbngWIFfvp6amul92tzW1dbV3weOIvd8+8mPEMeaAf4cEBAbA09ba9Wt1W3WR2yJhgI+Nj1V8CHOH6ImA3PsnfNsjxiDhFCQ6QIAk4ExLBIgAIQQQmRJEgBACiEwJIkAIAUSmBBEghAAiU4IIEEIAkSlBBAghgMiUIAKEEEBkShABQggg8l9ITjYwKjepmgAAAABJRU5ErkJggg==",
      "text/plain": []
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# second image in a batch\n",
    "ims[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# we'll use three operations\n",
    "from einops import rearrange, reduce, repeat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAGAAAABgCAIAAABt+uBvAAAHEUlEQVR4Ae2bDUxWVRjHI0C+FZAPFSS+CRrgApkSI0uRgROWLtSwtdVCEBkbzGw6P5pWw2JNECk2IcvICQgNCIgPISgVkYKhICAQCGIS8iHwIig99LYri/fe57z3nveFrcPeueN5/uc8z/nx3Pvec56Lxkxf33Psh5/A8/wmZpklwAAhecAAMUAIAcTMMogBQgggZpZBDBBCADGzDGKAEAKImWUQA4QQQMwsgxgghABiZhnEACEEEDPLIAYIIYCYWQYxQAgBxMwyiAFCCCBmlkEIIC3Erhbz/QcPWtrbu3p64PNweHh8YmJCJgPP+np6erq6ZqamL1hb265e/ZKLi6mxsVoieuZkwQANj4zkFRf/VFX1a10dcHkWkWDL2d7ed+3arQEBwRs36uroCGrpGDXUX1mtqKlJTk8vqqiYfPxY9CKMDA3Dtm6Nj4x0dXISPQnJQLUCKigtPZaYeKOxkSQyEo2GhsYbQUGfHzliZ2NDohehURMguMXEHj4MF5SIENEhcJ86FBv74b59mpqaqFhZgToApWRk7D9+XH7fVTY+cv3rfn7fnzljYWZGPoREqVpAssnJt2NisgsKSEKRrrGxsqrIynKwtZU+FTeDCp+D4As7YMcOtdGBJXX39r66bVtbZye3POkNVWUQPMtsCgu7cuOG9BCVncHR1ra2qMhk2TJlByrUqySDnj59+mZExILQgUW2d3XtioqamZlRuGBlO1UC6JOkpB/Ly5UNhaK+pLIy9dw5KhPSv8Su1de/Ehr65MkTKvGJnmSpkdGtqiqrFStEzyAfSDmD4OKKPnhwwenA2kZGRz9KTJRIB4ZTBvRtdjbFB2WJy/v64sU/7t6VOAlNQHBfTEhJIQ9IZ8kS2Ch8k5QE1wI6qr+h4Up+/rH4ePJHwampqaSzZ9GZhQU070HFly8HhYcL+5NbYQ/1fng4rHalpeW/PatWCQ/kNtWjjx5FHjiQmZsrrJdbgWZvfb2WlvhDC5oZdD4nhyRo2IjnZWR8dfIkR4dkFKeB4edPn353506uR6Dx58AA/NoEBKiJGiA4u/ihpAT1B/vJrLS0kM2bUaWAABLwy4QETzc3AQ1nKigr49oiGtQAwbf7o7ExNIL9UVGBGzagMlSgra2ddOIEKgOBxCMEaoB+vnoVDReeTT6IjkZlhAL/dev8fHxQcWd3d4+EP7egBqixuRmNdXtwMK0tktxX3J49qFMQ/H7zJolMoYYaoLaODoUO5nZu8vef+1/pbbhaSU6mGxYDoL7799EFv+zujmqUEkDZA87J0CGtBL88vkmoZRDJHXq5iQlfHKL7fdasQcf29vejGj4BNUAkJQrjpUv54hDd7+bsjI7tvXcP1fAJqAGCk3M+H1z/8Ogo16bVcHF0RKeCs01UwyegBghuB3w+uP6BwUGuTatBctlKqRdQAwQFYnTNTS0tqEZZgaG+PjpkUQAiKd2p4pjR0MAABSTl+JVaBtkT1DZzi4qk3A4UgiDJDgOCLFM4OXRSA7SW4Ot2aGTk41On+EIR1w9zogNJLkO+SagB8vX25vMxt/+LtDTIo7k9EttDBN9QJJchXxjUAEE908nOjs8N1w+H1m/t3Zt+4QLXI7Hx18OH6AxwhIRq+ATUAIGDsJAQPjdz+6Ee/V5cXOCuXWXV1dPT03NNItok34xQkhYxs3yI+LPI+S7fCQv7NDkZcmS+aX4PHNPAB5IfNmiW5ubzBYQ9dQQv07xI8DDJ545mBsEltn3LFj5PCvthBwcHSVn5+QqtJJ11DQ2obLEAgkCPxsXBWR8aMS0BvMfX3NaGzraIAMF7lvFkh1joqkgEcBdDi5RQXPIgO71W6JHmJSZ3cDQ+3tvTU6Ez6p35paXonOu9vUkO1fjmoQ8IoslNT19hYcHnklY/5A689IjO9pqvL6oRENAHBM6sV66szMmR/uKAQNxgKq+pIXkIIjlyFHCkEkDgz8XBoTovz93VVcC3RNN3ly6hM5gvX77eywuVCQhUBQhcwv7+WmFhxO7dAu6lmGB7DAVI+E0IVJZDAwMlvvpKszbPt9pfrl+POXTot6YmPgFJP1ebny+Gx/GO7u7bd+7cbm+f/fefDxSdQVmcmSmxTqkOQBAonMgUlpV9lppKUl+cj2B2BiWLf7DLB15eHh4C+aXQ0X861QSI8wovEF7Iy4MNPRTzCDcl8rHKAuI8SmyoGxAXLjwE19TWwttWt1pb4WkYKjODQ0MCR3//O0AcKa4B7zs9GBwcGx+XyWTyItIS+NHWhvNA+CsoKaeCnAsRjQXLIBGxLsgQFX7NL8h6qDtlgBCkDBADhBBAzCyDGCCEAGJmGcQAIQQQM8sgBgghgJhZBjFACAHEzDKIAUIIIGaWQQwQQgAxswxigBACiJllEAOEEEDMLIMYIIQAYmYZhAD6G5tIOOckeJAKAAAAAElFTkSuQmCC",
      "text/plain": []
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# rearrange, as the name suggests, rearranges elements\n",
    "# below we swapped height and width.\n",
    "# In other words, transposed first two axes (dimensions)\n",
    "rearrange(ims[0], \"h w c -> w h c\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAGAAAABgCAIAAABt+uBvAAAHEUlEQVR4Ae2bDUxWVRjHI0C+FZAPFSS+CRrgApkSI0uRgROWLtSwtdVCEBkbzGw6P5pWw2JNECk2IcvICQgNCIgPISgVkYKhICAQCGIS8iHwIig99LYri/fe57z3nveFrcPeueN5/uc8z/nx3Pvec56Lxkxf33Psh5/A8/wmZpklwAAhecAAMUAIAcTMMogBQgggZpZBDBBCADGzDGKAEAKImWUQA4QQQMwsgxgghABiZhnEACEEEDPLIAYIIYCYWQYxQAgBxMwyiAFCCCBmlkEIIC3Erhbz/QcPWtrbu3p64PNweHh8YmJCJgPP+np6erq6ZqamL1hb265e/ZKLi6mxsVoieuZkwQANj4zkFRf/VFX1a10dcHkWkWDL2d7ed+3arQEBwRs36uroCGrpGDXUX1mtqKlJTk8vqqiYfPxY9CKMDA3Dtm6Nj4x0dXISPQnJQLUCKigtPZaYeKOxkSQyEo2GhsYbQUGfHzliZ2NDohehURMguMXEHj4MF5SIENEhcJ86FBv74b59mpqaqFhZgToApWRk7D9+XH7fVTY+cv3rfn7fnzljYWZGPoREqVpAssnJt2NisgsKSEKRrrGxsqrIynKwtZU+FTeDCp+D4As7YMcOtdGBJXX39r66bVtbZye3POkNVWUQPMtsCgu7cuOG9BCVncHR1ra2qMhk2TJlByrUqySDnj59+mZExILQgUW2d3XtioqamZlRuGBlO1UC6JOkpB/Ly5UNhaK+pLIy9dw5KhPSv8Su1de/Ehr65MkTKvGJnmSpkdGtqiqrFStEzyAfSDmD4OKKPnhwwenA2kZGRz9KTJRIB4ZTBvRtdjbFB2WJy/v64sU/7t6VOAlNQHBfTEhJIQ9IZ8kS2Ch8k5QE1wI6qr+h4Up+/rH4ePJHwampqaSzZ9GZhQU070HFly8HhYcL+5NbYQ/1fng4rHalpeW/PatWCQ/kNtWjjx5FHjiQmZsrrJdbgWZvfb2WlvhDC5oZdD4nhyRo2IjnZWR8dfIkR4dkFKeB4edPn353506uR6Dx58AA/NoEBKiJGiA4u/ihpAT1B/vJrLS0kM2bUaWAABLwy4QETzc3AQ1nKigr49oiGtQAwbf7o7ExNIL9UVGBGzagMlSgra2ddOIEKgOBxCMEaoB+vnoVDReeTT6IjkZlhAL/dev8fHxQcWd3d4+EP7egBqixuRmNdXtwMK0tktxX3J49qFMQ/H7zJolMoYYaoLaODoUO5nZu8vef+1/pbbhaSU6mGxYDoL7799EFv+zujmqUEkDZA87J0CGtBL88vkmoZRDJHXq5iQlfHKL7fdasQcf29vejGj4BNUAkJQrjpUv54hDd7+bsjI7tvXcP1fAJqAGCk3M+H1z/8Ogo16bVcHF0RKeCs01UwyegBghuB3w+uP6BwUGuTatBctlKqRdQAwQFYnTNTS0tqEZZgaG+PjpkUQAiKd2p4pjR0MAABSTl+JVaBtkT1DZzi4qk3A4UgiDJDgOCLFM4OXRSA7SW4Ot2aGTk41On+EIR1w9zogNJLkO+SagB8vX25vMxt/+LtDTIo7k9EttDBN9QJJchXxjUAEE908nOjs8N1w+H1m/t3Zt+4QLXI7Hx18OH6AxwhIRq+ATUAIGDsJAQPjdz+6Ee/V5cXOCuXWXV1dPT03NNItok34xQkhYxs3yI+LPI+S7fCQv7NDkZcmS+aX4PHNPAB5IfNmiW5ubzBYQ9dQQv07xI8DDJ545mBsEltn3LFj5PCvthBwcHSVn5+QqtJJ11DQ2obLEAgkCPxsXBWR8aMS0BvMfX3NaGzraIAMF7lvFkh1joqkgEcBdDi5RQXPIgO71W6JHmJSZ3cDQ+3tvTU6Ez6p35paXonOu9vUkO1fjmoQ8IoslNT19hYcHnklY/5A689IjO9pqvL6oRENAHBM6sV66szMmR/uKAQNxgKq+pIXkIIjlyFHCkEkDgz8XBoTovz93VVcC3RNN3ly6hM5gvX77eywuVCQhUBQhcwv7+WmFhxO7dAu6lmGB7DAVI+E0IVJZDAwMlvvpKszbPt9pfrl+POXTot6YmPgFJP1ebny+Gx/GO7u7bd+7cbm+f/fefDxSdQVmcmSmxTqkOQBAonMgUlpV9lppKUl+cj2B2BiWLf7DLB15eHh4C+aXQ0X861QSI8wovEF7Iy4MNPRTzCDcl8rHKAuI8SmyoGxAXLjwE19TWwttWt1pb4WkYKjODQ0MCR3//O0AcKa4B7zs9GBwcGx+XyWTyItIS+NHWhvNA+CsoKaeCnAsRjQXLIBGxLsgQFX7NL8h6qDtlgBCkDBADhBBAzCyDGCCEAGJmGcQAIQQQM8sgBgghgJhZBjFACAHEzDKIAUIIIGaWQQwQQgAxswxigBACiJllEAOEEEDMLIMYIIQAYmYZhAD6G5tIOOckeJAKAAAAAElFTkSuQmCC",
      "text/plain": []
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# we could use more verbose names for axes, and result is the same:\n",
    "rearrange(ims[0], \"height width color -> width height color\")\n",
    "# when you operate on same set of axes many times,\n",
    "# you usually come up with short names.\n",
    "# That's what we do throughout tutorial - we'll use b (for batch), h, w, and c"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Composition of axes\n",
    "transposition is very common and useful, but let's move to other capabilities provided by einops"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 11904 characters]",
      "text/plain": []
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# einops allows seamlessly composing batch and height to a new height dimension\n",
    "# We just rendered all images by collapsing to 3d tensor!\n",
    "rearrange(ims, \"b h w c -> (b h) w c\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 11068 characters]",
      "text/plain": []
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# or compose a new dimension of batch and width\n",
    "rearrange(ims, \"b h w c -> h (b w) c\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(96, 576, 3)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# resulting dimensions are computed very simply\n",
    "# length of newly composed axis is a product of components\n",
    "# [6, 96, 96, 3] -> [96, (6 * 96), 3]\n",
    "rearrange(ims, \"b h w c -> h (b w) c\").shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(165888,)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# we can compose more than two axes.\n",
    "# let's flatten 4d array into 1d, resulting array has as many elements as the original\n",
    "rearrange(ims, \"b h w c -> (b h w c)\").shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Decomposition of axis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2, 3, 96, 96, 3)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# decomposition is the inverse process - represent an axis as a combination of new axes\n",
    "# several decompositions possible, so b1=2 is to decompose 6 to b1=2 and b2=3\n",
    "rearrange(ims, \"(b1 b2) h w c -> b1 b2 h w c \", b1=2).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 11200 characters]",
      "text/plain": []
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# finally, combine composition and decomposition:\n",
    "rearrange(ims, \"(b1 b2) h w c -> (b1 h) (b2 w) c \", b1=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 11544 characters]",
      "text/plain": []
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# slightly different composition: b1 is merged with width, b2 with height\n",
    "# ... so letters are ordered by w then by h\n",
    "rearrange(ims, \"(b1 b2) h w c -> (b2 h) (b1 w) c \", b1=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 15612 characters]",
      "text/plain": []
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# move part of width dimension to height.\n",
    "# we should call this width-to-height as image width shrunk by 2 and height doubled.\n",
    "# but all pixels are the same!\n",
    "# Can you write reverse operation (height-to-width)?\n",
    "rearrange(ims, \"b h (w w2) c -> (h w2) (b w) c\", w2=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Order of axes matters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 11068 characters]",
      "text/plain": []
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# compare with the next example\n",
    "rearrange(ims, \"b h w c -> h (b w) c\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 13184 characters]",
      "text/plain": []
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# order of axes in composition is different\n",
    "# rule is just as for digits in the number: leftmost digit is the most significant,\n",
    "# while neighboring numbers differ in the rightmost axis.\n",
    "\n",
    "# you can also think of this as lexicographic sort\n",
    "rearrange(ims, \"b h w c -> h (w b) c\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 11096 characters]",
      "text/plain": []
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# what if b1 and b2 are reordered before composing to width?\n",
    "rearrange(ims, \"(b1 b2) h w c -> h (b1 b2 w) c \", b1=2)  # produces 'einops'\n",
    "rearrange(ims, \"(b1 b2) h w c -> h (b2 b1 w) c \", b1=2)  # produces 'eoipns'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Meet einops.reduce\n",
    "\n",
    "In einops-land you don't need to guess what happened\n",
    "```python\n",
    "x.mean(-1)\n",
    "```\n",
    "Because you write what the operation does\n",
    "```python\n",
    "reduce(x, 'b h w c -> b h w', 'mean')\n",
    "```\n",
    "\n",
    "if axis is not present in the output — you guessed it — axis was reduced."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 4832 characters]",
      "text/plain": []
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# average over batch\n",
    "reduce(ims, \"b h w c -> h w c\", \"mean\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 4832 characters]",
      "text/plain": []
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# the previous is identical to familiar:\n",
    "ims.mean(axis=0)\n",
    "# but is so much more readable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAGAAAABgCAAAAADH8yjkAAADw0lEQVR4Ae2ZT0hUURTGz5vUHFEhrcxpsFIrQoS0IM0/ERS6qZWQ5KJlQW2DdiLYojZF0KbaRFkEQYuoXSFqUZAgQYYWkaXVGKlhjiNjU9b43tx73v3uu+OMu3Hzzvfdc77fvDtv3jwZa5BW98+3uvFEWYaAX3f6xyPF1a0HDPudNstsi67cisVHNl8ud2aNCrMtOnlzOZ8m2vuNcp0mI0D3K6effp8dTwiDygQQeiAGRTtF5VmbAHrklKEfstYrE8AzFtHHtFaaACZYQlJvggkgygAzTGulCYB/GAu1iWzRBFDKZoJMa6UJoJElNDGtlSaA43JC9UZZ65UJIHBUzMjqEpVnbQKgztpEju/iloQwqIwAdL3D7iu9m+QN2/B2TbM9AxPzxdUtBw1etNRiCpCGkhH2qSczk1RvBuC5XfxGphpYGPs6Gfoeml5YiEQWrT9ZuX5/XmkwGKwsUnUzD19Fe1inQgZq6/d73VpTAiwxsxva6iwF27FSvYqivWceOmGqIlXAUmaZKtfx0gCocMJUReqAkgJVruOlDihyspRF6oC3l5TBtpk6gG6P2GGqYxoA1B1TJS976QAMP1oRYHBwn2ZOWrrGn/2EVc0ZvH4p9Ellm6SIvjxmhiA1gHtCm1yeaJY13WdakBjw84nQxsoL7P42jC8kDOjVbGxOAyM+ZTohNYBEk7s6xKw+phMSAmLCP36Jdrvi30ajM/YKP0LAhzBvFXVAFP/qYW7YGgLe2B3q4zpmQwD80v/EEkR5RBTx+rPbijvwDL6hCbUP2yEgpA5C7iRagADte+xOm3dbcQcCImhC7cN2CFhUByEXfuwhYC2KUvuwPV2AXDWXCALy0YTah+0QUKIOQi5sh4BNKErtw3YI2KYOQu5WtAABu9DEkr/DvVbltuIOBJTBt43o9Hn+PLpmZ9IAqw6NEI203mDX/W4/6oZnQE1ohOg5VbbIqzWyFBQGNOcIbXI5NBebkp29shQU/MKhwsP4ibBrw4CQQZSPzwAD6BgG8EemehyDt4iq+MOP9KIl8UJSktAA6JTUqROzeFEHwFOulXcuxzZ0gGm7yfs4Clt0APyyXHG4NU2AlZ3Be9cLhcaKAItjMM+1gN8uzRaFFXdlV7JtfLQLftQACs/xZqyz59Aa/oyjCYUfaO9QuHFLcwZwxrXQiPONfwkk2j4ViUZ9vpifwlZeXnjeyi9YX15RWWO5eJJhvkVXi6VBU5GWLdLBMgDd7vxfy2xRZos8d8CzIXMVZbbIcwc8G/BvOJ6jZg2rfpn+BUFLlGkzaXw/AAAAAElFTkSuQmCC",
      "text/plain": []
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Example of reducing of several axes\n",
    "# besides mean, there are also min, max, sum, prod\n",
    "reduce(ims, \"b h w c -> h w\", \"min\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 5480 characters]",
      "text/plain": []
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# this is mean-pooling with 2x2 kernel\n",
    "# image is split into 2x2 patches, each patch is averaged\n",
    "reduce(ims, \"b (h h2) (w w2) c -> h (b w) c\", \"mean\", h2=2, w2=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAASAAAAAwCAIAAAA+QDD4AAALi0lEQVR4Ae1dCVAUVxoWZjw4hssVOQXxigVqdILHBI2uCq6skpJ4FRo1RkLWFEaT0i2TRa21YjxILNm4uOWuujFKFBKDF2LwWG4FAYsFFRUvGC1QkEMCI7B/ZJ0M73UP73XP0Q4vRVXe/73/el/3P9393uvWqr2ysoeU/lO7SykbyEViCUksnR5jrkjseEksIWuJ0SMqnU2fbHrN+bWlYUs1LRpRjpgxY8BADMgN5MfMbn5p+sXP1q8jidTkVJ/ePjv375y3ZJ6Z02Lhuz0DFnIF01aX9oB+vPRjbZs1GAPmYsBCCoyTvt3bd3PiDGQMmIwBSyiw9vZ2Tr7u3b7HiTOQMWAyBiyhwKysrDj5Wrh8ISfOQMaAyRiwhAIDsmbPn41QJpfLR70xCgGZyBgwMQMWUmDxCfERKyK03Ll7ud/TsPtDLR+sYTYGrNhCcxfcS2xlV2LpsIXmLk4fC7mCdTFK1s0YMBMDrMDMRDwL2z0YELuT4z85OX/bt+9idvazpqYBnp5Tg4JWvf/+IF/f7sEe7ygfParMyrqQm5teVJRXU/MY/kDVy8tHpZqyfHm0r+8gXksjdyTv25dy+HBpfn57W5ufv//U8PCFq1ZZW5v6d7boRlFmUWZucW7BtYIndU8e1z5W2ClGDxsdNjksck6kXCb2tBTM4sGDp3744Vxh4Y2HD6ttbHq7urr4+w8KDh4fFja5f38XAW4FPoO1tbUFTJlSWlbGF3LurFlH9uzh69WDkzxjqB+old5KPU60XZXtorcykyT0Il5c3Bb4a2io10bX03j33agvv/y7HgW+LuJ0Ojm4XVIyf+TIttbWTqiOsDs1ddz06ToAaZN8b+3mvZs37tnY2sabg27I+PXxH4R/oIuQtskT0vG4Z09SVNQXOgB3s1evnsnJX4eETODu5kKFFNjmnTv/sm0blzcUu5+f7+VOtz2e5ASSZoF5eHAvx6Gk6MiVldxL5DoqaJOEH8Rmtp9fRXk5AuKirUKRXleH4/oR8vPZSknHDyxvPr/0nPrqSp7Qi4HBpUIuH8u3V4Fz7O3t+Zw4J0h9b6CaNYuwuiCet1JZfu8eZ2AGAgMCapKWt0CZjKS6wO2z+nolz5I9bVCD6MNJLwuUNTU3GcQbnxOZLJCqumbMUPG54sTpCmxuZGR2PkX5Qki/8eM5AzOwg4GoqAXGo0Jlawu/0FT+JVVjkLmtypYqfyplB4eJVPqgvGBBCJUJRYFdu3kz8cQJKu8dymOCgwVYdROT5OTvjTTSL6KimpuE/Py/9+abRkpJmNvAxYHCDPVbNTe31Nc/06+D986bR/ekSlFgwydNQuLJZLInpaWwVK3790lUFKJWUFyMICJF2KjRKeSL8MVVBo4iMkkwLyxUd+am/Z13FuNu09JO4aB4JIlrkimloiIfniFe/l3SaKxlMiRWUVaWpqUFAY0hqlPVv6XyIqfEbYl4oLySPMKpEdxWDxITE4/31tenv+RGS9Kvjby8gyNGDAZ9mFrErfQgpAV2LiMD8fLrM+j9+86Ojgi+Iybmn199hYAwL4Ig3VPctevf+NZkmHs0OBufRfy2cUzrPK+trZ+Hh1aEhkwuv/z8uS7S0TbXRSx8avit5Ft4Pm+veRsHRSIpKVmIh4SELfb23HekSuXwq1e/h0pDTLoUSQts5qJFiK+KK1cQRCu+twB9rtiwY4e2t5s3NmyIRRi4dAn98UIUBIgphw4hVqtjY/Ha7tCJPXYMUS7Jy0MQk4l+nn5DfYYi4U6kn0AQ8SJeSzdv3hfvFvFAWmDN2D2De//+iC9dcVLnuQ3aR21dVxbWhhUwY4+Ik+1Fa9bwxZ0cFoZ31VZX46BpkPyD1BcKAYktWfJHxOrzz2E5MAcBRYpEBfaoqoo2TP9+/WhNuol+nz42xh7p8f37xYdIjOd4RBHvlsSDva09rpZ9NRsHxSCRkXNw85CQlVZWyvj4RLxLGEJUYCnnz+PerWARh//v6PHjiIlGo0EQJhqJgUs//4x49g/sYiJuROc7DjC/fO4c4sS8YtZV9JFJfD4//YROFnT4/PDDLVBmcnng/v3oaUwblKjADDINeOfBA9rkmL4wBsqvXUMM3QYMQBBEdPX0RJC7168jiClFFwcXJNyt+xyTH4gOrTh79ltxcWv5rFpb25Yt2wiV5uo67dGjJ3xq+nGiAoPNqvq9kPQ+pd+GQ+KW6eAMaJqbEbC3TRf3pbhCC+YE8WlU0VGBzk4baUvHRx/Nr6tL79lTrmc4VVU1bm7Tocz06PB1ERVYX2dnPnty3N7OjlyZaYphoK+bG2Je09VTdBX2geffYU4Qn0YV6xvrEf8ujug1DVEQLCoUti0tuTU1F1SqUXqcQJnB1ay2Fk1Mjwl0ERWYcuRI/V5Ier07r8CQmFDpwEc4qPQtWHmUSoWM7r+XLyMIIhbn5iLIKLPu56iurUbyCRgUgCCGFZ2cFJmZ/4KVroqKlGnTxvE5d3aezNfFiROdlKHTOC6OUvvWgJ2CXSH/f4hh0XLv5s26x7vuSRePEE2Njbr60A5dvBhBzCsGjw82TQIeHv3Ont0NsWAj1cSJy4uKbiBx585dd/ToVgTkE4muYC5OTnz20sFh35Z0kjFvJj7DhuEJ6NkzDvvocf3Xg4Jw0IyIpys6DWPsZODWsbDwMGzgQAIlJqKTtIiCrkhUYGCgsLfXNYN2xMqVCMJE6TDQq08fJJkIpRJBtOIMLy9tu6NB/RYWYi9ODI0OFefAkNYdWxAFeyQtsNyTJ5EYh3788XWuW0dELfXiRQQxpbg4VFr3OSYb+16M9usFBZmnT+MJbI+ObsQmeL9OTsY1TYPcqbxzKvMUEiviDxEIIl6E1WRr6zcMuKbMmRJpgQ0fMsQG+1EsKimBtWabgQM3xcZeLS2FifjW1tYHajV8omPZ6tVyb2/oDVm4kDOwMUAfPx/EbdqpNNVgVUlRCSTW2NCYn5O/Zf2WIYohwxw5bqIQ21da9B87Fr8KRc+cOcPTM+/8eXiNpeHp07SkpAk2NglxcfhIg0JNcQ1xD3b3n+ufcCYBpjQ0zzUlt0um/2n6wFkD8Xy+/eu3OCgS+e6703Db3LGmDNODERGfnTyZUV1dCyCsgD1+/DQzs2jdul3QJSYQxScD4ByFmhEQ7HpGxlA/P0JDAa/Eaz0XFxQHjyF9FCb9XAdxQvjryfC6iqsrOmMO2eKa5N8OIE6nB8xbBGE39lqu9DTOPnzoonejqa4t+Rv6tJ8M0EYZ7D247FiZVuyiQZyQ4MoJD5+amLitizRedpNewUAfZhFuZmW9NKT4f/KZMxTaIlQDRgeIsLY0Uxs7u+1JSbSjgk335NVF61yYPkV1CQtAaUVeXeCYosBAG77HBm9YUubTIzk1ldZEsP61WnSXkGBXFmD4+zlztiUmkg8keutWPZvuyf0YULP8eLkBvYl3pVbTncx0BQb5wRuWsAK2gut9Pr7s07FFTD5N8biDo8OVB1dI/GRfzCZRe9V14MuHmY2N+GvL+LhS1eola9fiuBmRmgs1vh6+RkogJ+eAnZ0NufMBA9xgGdrNrS+5CWhSF1iH939s3w5ldvzAgYH8u0h79+q16dNP4a1nEy9Ju3m6wfNV9PpoTiK8fb2PpB0BhQlvTeBUsDywj60tvLZ8rKyMc8uvg7PzgZwceC0e32BlbCrgkwEVKRUjh3DsE/rmz9/A1wScFE7Gy2HcuICGhgyomR07VuvZiyiTWcfErAC1u3fRiXSS3CgmOUjcidchf4gXH4vIg8QSklg6FP/4Az7JAQXm1pdjEojouPApEU9y8DkwLC7wCmbYJJg3xoClMsAKzFKPLBuXJBhgBSaJw8CSsFQGWIFZ6pFl45IEA6zAJHEYWBKWygArMEs9smxckmCAFZgkDgNLwlIZYAVmqUeWjUsSDBB9MkASmbIkXnEGYGfGKz4CIen/D3pQsiIrI2XcAAAAAElFTkSuQmCC",
      "text/plain": []
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# max-pooling is similar\n",
    "# result is not as smooth as for mean-pooling\n",
    "reduce(ims, \"b (h h2) (w w2) c -> h (b w) c\", \"max\", h2=2, w2=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 5992 characters]",
      "text/plain": []
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# yet another example. Can you compute result shape?\n",
    "reduce(ims, \"(b1 b2) h w c -> (b2 h) (b1 w)\", \"mean\", b1=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "source": [
    "## Stack and concatenate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'list'> with 6 tensors of shape (96, 96, 3)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(6, 96, 96, 3)"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# rearrange can also take care of lists of arrays with the same shape\n",
    "x = list(ims)\n",
    "print(type(x), \"with\", len(x), \"tensors of shape\", x[0].shape)\n",
    "# that's how we can stack inputs\n",
    "# \"list axis\" becomes first (\"b\" in this case), and we left it there\n",
    "rearrange(x, \"b h w c -> b h w c\").shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(96, 96, 3, 6)"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# but new axis can appear in the other place:\n",
    "rearrange(x, \"b h w c -> h w c b\").shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# that's equivalent to numpy stacking, but written more explicitly\n",
    "numpy.array_equal(rearrange(x, \"b h w c -> h w c b\"), numpy.stack(x, axis=3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(96, 576, 3)"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# ... or we can concatenate along axes\n",
    "rearrange(x, \"b h w c -> h (b w) c\").shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# which is equivalent to concatenation\n",
    "numpy.array_equal(rearrange(x, \"b h w c -> h (b w) c\"), numpy.concatenate(x, axis=1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Addition or removal of axes\n",
    "\n",
    "You can write 1 to create a new axis of length 1. Similarly you can remove such axis.\n",
    "\n",
    "There is also a synonym `()` that you can use. That's a composition of zero axes and it also has a unit length."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(6, 1, 96, 96, 1, 3)\n",
      "(6, 96, 96, 3)\n"
     ]
    }
   ],
   "source": [
    "x = rearrange(ims, \"b h w c -> b 1 h w 1 c\")  # functionality of numpy.expand_dims\n",
    "print(x.shape)\n",
    "print(rearrange(x, \"b 1 h w 1 c -> b h w c\").shape)  # functionality of numpy.squeeze"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 10664 characters]",
      "text/plain": []
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# compute max in each image individually, then show a difference\n",
    "x = reduce(ims, \"b h w c -> b () () c\", \"max\") - ims\n",
    "rearrange(x, \"b h w c -> h (b w) c\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Repeating elements\n",
    "\n",
    "Third operation we introduce is `repeat`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(96, 5, 96, 3)"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# repeat along a new axis. New axis can be placed anywhere\n",
    "repeat(ims[0], \"h w c -> h new_axis w c\", new_axis=5).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(96, 5, 96, 3)"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# shortcut\n",
    "repeat(ims[0], \"h w c -> h 5 w c\").shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAASAAAABgCAIAAAAyzjgPAAAJXElEQVR4Ae3de0wUVxQG8KKCivhWxKgUFKWiVKtohBCltaL4ahpbNNEG01ipr9BIrI2Nj/5RjUlJoxXTaESNmJrIKqYiUhBFrEUFW6gKitLWB2q1KL7f9DSbbDasnBl37p3duf0a0wxz7txZf3c+Z9mZvePTUFv7Gv6DAATkCDST0y16hQAE/hNAwHAcQECiAAImERddQwABwzEAAYkCCJhEXHQNAQQMxwAEJAogYBJx0TUEEDAcAxCQKICAScRF1xBAwHAMQECiAAImERddQwABwzEAAYkCCJhEXHQNAQQMxwAEJAogYBJx0TUEEDAcAxCQKICAScRF1xBAwHAMQECiAAImERddQwABwzEAAYkCCJhEXHQNAQQMxwAEJAogYBJx0TUEWniQ4PGTJ8dOnjxcUlJRWVldU1N7/fq9+/dpZetWrfxbt+7SqVNocHDv4OBhgwfHREX1CQnx4Ev1yK7hw7NbwsfH/IlHGxoa8g4dyrTZ9uTlUaJ4REe1b2ho4uTJSYmJtOBYqeQCfPhhtZaPqQF78eLFtqys1enpldXVPGJT1WbNmk2ZMGH5woUDwsObamPd9fDhx86KPuYFjN4NzluypKyigkfUU/X19U1NTl6emtqqZUs97S3RBj78MFnUx4yA0T88K9euXZGW9vz5cx7xlapRgwbtzsjo2b37K23lhY3hww+KpX2kB+zBw4cfzp6978ABHtG9alBg4CGbLbxPH/c294at4MOPgtV95AbsVn39hBkzfikr4xGNVHsEBRVnZ9PnjUY68dS28OHlFfCRGLBHjx+PmTr1yPHjPKLxamT//sdycujDfeNdmdkDfHhtNXwkXmj+aMECE9JFg/R7ZeVny5bxo+WFVfjwg6KGj6yApW/enLV3Ly8osLohM/PnEycEdii7K/jwwsr4SHmLWHX+/JD4+IePHvGIYqtvDRxYlpfn4+MjtlsZvcGHV1XJR8oZLGXpUpPTRQP266lTOQUF/Mh5SRU+/ECo5CP+DLY3P39SUhIv6Fxt17btlPHj3x05ckhkZOeOHTu0a1d/9+7NurpTVVX04f7u3Nzbd+44t2eWR44YUbRrF9PAG0rw4UdBMR/xAYsaN07n7RrNmzdfNGfO5/PmdWzfvil0+qD26zVrvt2wga42NtXGeX310aNh3n1bMHycx8t1WTEfwW8RC48c0ZmutgEBOdu2rVqyhEkX6VP1m2XLsjZu1HlX1I7sbNcx85418OHHQj0fwQH7LiODF7RX6aOIzHXrxsbF6WlMbd5PSEhftUpPY3pLqaeZp9rAh5dXz0dkwOrv3MktLOQF7dVPpk+fHB+vp6WjzcfTpsWPGuX4samF306fppfRVNWz6+HD+yvpIzJg2fv303fgeESqtvTzW5GaqtnMtcGiuXNdVzZaQ7+qmXN1u9F+9fwIH15JSR+RAfupqIgXtFfHjx7dvVs3PS0btYmLjg5o06bRStcfdf4S6Lqh7DXw4YWV9BEZsKOlpbygvUq/UOlp5tqmRYsW9FG+6/pGa86cO9dojZf8CB9+IJT0ERaw6zdu/HnpEi9or9L3uPQ0e2mbbl27vnS980q3vy7t3InwZfjwpKr6CJv0hm5v4QUd1Qgdn1U4GruxcOXaNTe2kr0JfHhhVX2EncF0nr54ZSHVutu3nz59KqQrgZ3Ah8dU1UfBgNGsQzfq6vjhNL/qPQcQfPjRF+sjLGB0TxP/us2s3n/wwMzd6dkXfHglVX2EBYzmTuAFzaw+MvebMnr+avDhlVT1ERYw87+fwgyYnuvdzOYySvDhVVX1ERYwng9VCPw/BYQFjGaT9x5BPz8/73kx9lcCH35EVPURFjCvmtTJz9eXH07zq/DhzVX1ERYwehgKL2hmtY2/v5m707Mv+PBKqvoIC9jrPXvygmZWO3XoYObu9OwLPrySqj7CbpUK6dWLF3RUr5WX67ml0NFejQX48OOoqo+wM5j+5wn9cfEib61kFT78sKrqIyxg9K6sX+/ePKK9Sk/f09NMsTbw4QdUVR9hASO+mGHDeER7df3WrXfv3dPTUrE28OEHVEkfkQGbNGYML2iv/n3z5qeLF9MtlXoaq9QGPvxoKukjcl5EehxGYGSkzrMTTWLz/erV9KxKHl2lKnz40VTSR+QZjKYuTJw0iUd0VDN27BiWkHC4pMSxxo0FemQmTeSQlJLyVVqaG5ubvAl8eHAlfUSewYiPvq4/IC7uld7+xQ4fvjA5meZI1H+zDM3vVVBc/GN+Pk2z/M+tW7RfmgRuz5Yt/Ph5QxU+/Cio5yM4YMQ3ZdasXfv28Y6uVfrX653Y2OGDB0f06xceFkaT1Af4+9McUnSTNc1Nf7u+noJEs9WXVlSUlpfTMDR63DM9RbaquNi1Wy9cAx9+UBTzER8wusxFJzGTv31AE049rKmh//OD5w1V+PCjoJiPyN/B7HD0uOQvU1J4ROHVZ8+e1Vjk+jV8+NFXzEd8wIjvi/nz6f0e7yi8evbCBeF9SuoQPjysSj5SAkbPJfph/frgHj14R7HVs7rnjRO7Xzd6gw+PppKPlIARX2CXLoU7d/YICuIpBVYtdAaDj+a4K3P8yAoYCfYJCTlos5n2ODxrBQw+mhlT4/iRGDAS7Bsaejw3V/9zwDTRmQaWCxh8mNG0lxQ4fuQGjJjoEZW527enr1xJz2LWBDXSgG5x1P80ZyM7ErstfHhPq/tIDxjx0fMs586ceaaoiJ67J/XmQwt9zuF8VMHHWcN12dI+4i80uwI5r/nr8uW1mzZl2mx0wnFeb2S5a+fO740d+8HEiaNjYy1xrZn5y8KHwaGS5XzMDpidj64L7z94cG9BAd2q694XnOkxmdFRUW/HxNAFt+ihQ+mDXX5grFWFDz9eFvLxTMCc+S7V1tKDlctPnz5XU0NPHrpy9SpNU053WtEfummY5oey35TYNiCALqy9ERZm//NmRATdvujcj6rL8OFH1st9PB8wng9VCFhawIwPOSwNhBcPASMCCJgRPWwLAQ0BBEwDCGUIGBFAwIzoYVsIaAggYBpAKEPAiAACZkQP20JAQwAB0wBCGQJGBBAwI3rYFgIaAgiYBhDKEDAigIAZ0cO2ENAQQMA0gFCGgBEBBMyIHraFgIYAAqYBhDIEjAggYEb0sC0ENAQQMA0glCFgRAABM6KHbSGgIYCAaQChDAEjAgiYET1sCwENgX8B86gmBRxzD24AAAAASUVORK5CYII=",
      "text/plain": []
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# repeat along w (existing axis)\n",
    "repeat(ims[0], \"h w c -> h (repeat w) c\", repeat=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 5480 characters]",
      "text/plain": []
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# repeat along two existing axes\n",
    "repeat(ims[0], \"h w c -> (2 h) (2 w) c\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAASAAAABgCAIAAAAyzjgPAAAJxUlEQVR4Ae2deWwWRRjGbYFy37dA5RKUWywISKAKiFLEIIqCJEQkEIUEhUgUNGI8CCaNgQAeCYgRgwpFiAVByyWIQECFcF8iRwHBclNu1DzMG0qLu939Zna/2ad/fHm/nZn3nfnNPN1vd2dnEq5nZ9/BPxIgAT0EEvW4pVcSIIH/CFBgHAckoJEABaYRLl2TAAXGMUACGglQYBrh0jUJUGAcAySgkQAFphEuXZMABcYxQAIaCVBgGuHSNQlQYBwDJKCRAAWmES5dkwAFxjFAAhoJUGAa4dI1CVBgHAMkoJEABaYRLl2TAAXGMUACGglQYBrh0jUJUGAcAySgkQAFphEuXZMABcYxQAIaCVBgGuHSNQlQYBwDJKCRAAWmES5dk0BRIjBM4OKlS4i49tdfYfy0Zg2MTdu2wdi1dy+M7KNHYZw9dw6GFC9ZogSOlCpZEkaVSpVg1EtOhlFfGW1atcKRDikpMBrUrQuDn1oJJHDhUR18r1+/DreLly+HMTMjA8b8xYthiGbw1fDn3fXqIWLfXr1gDOzbF4YkGa6SleEosBh067Vr1+DlizlzYEyYMgXGtl27YhDAiIvExBvXC33S0hDwrZEjYTRt3NhIFSwMQoF571T5jTdszBh42bBpk3d34StZrFgxVGrU0KEw3ho1CkaJ4sXDV98w1ogCc9srcpp6f9IklBmXng7j6tWrbr3Eeb6Uli3Rgm+nT4dRu2bNOG+T3upTYA58z+fmIsfTQ4bAWLhkiUOZCCTXqFYNrVyuri0bN2gQgXYXuokUWMHITpw6hYS0AQNg/LJhQ8FZo320Vo0aALBy3jwYcg8z2mButJ4CyzMMLly8iO/dnnkGxqp16/Lk4JfbEGh+771IWbtgAQx5kHCbEpE4TIHl6Wb5HTgnMzNPAr+4JjBEnfM/+eAD14WszUiB/de1Uz77DD08fOxYa7vaeMNWzZ+PmA+2aWM8eFgCRlpg23fvRj+0fuQRGLkXLoSlZ+K/Hvc1a4ZGbFDP1hMSEuK/WYVrQaSnSo14803Qoq4KN2rc5f5t82ZkXJCVBaNnt27uitqTK4pnsMwff0QHPj5wYBh6slzZsqhGnx49YHTt1AlG6+bNYVSuWBFGhXLlYJw6cwbG8ZwcGJu3b4chDxK+/f57HDl5+jSMQD47tWuHuCvmzg2kAgEGjaLAUh59FMQDmXhRpEgRRH/1xRdhjB42DEbF8uVhxOpTHja8N3EifH746acw5Ll5rGK58bNr9WpkaxiZqcYREtjSVavQu13UrFY3YyJWecqWKQNXs9UQ756aGivn7v3IOa3/Sy+hlDyZcO/Ec853Ro9G2Tdeftmzk/gqGCGB9R40CH0zb9EiY50kl/Xz1I3KXuqGirE6FBho+ldf4fgLakZvgdlie1B+8cptj9j6D6E3+wV2Sl1+VG/RAh0gr1QZ6I/wPxTq3q8fOPywYoVuIDJhP2frVsQqry4pdYcOyr/9dxHlfGVSV8WTktCj49T086A62DHuq+q3ogGByYWfzI9J69rVsYZxncF+gRkYN/lHQI8uXXCwZvXq+VNDdSS1fXvUp0zp0jAMvAkqt5cosFANBi+VWb1+vZdi/sr0fuwxfw7MlS5a9MY/WblAkiUM9FVi686d+pyHyrO1Z7Cjx44B9L4DB8wTl/emzIf2HLF61aqeyxa2YBy96F3Ypt2S31qByTSoWxps5muTzp3NBIrTKIeOHInTmhe22tYKLJATV2HpRzZ/zsmTaPvly5dhyPIEljGhwCzr0Phojqy6dUzN87oz9HeDvJG1VmAyS8gbF5YyQ+Dc+fNmAgUVxVqByVoaQZFlXDcELtj+fpC1AuMbKG7Gd+B5TD79D6SxiYFEZVASiAgBa89gsmJ7RDoyTpuZpOaUxWn9HattrcC4pJFj34chQ5JaPDgMldFRB2sFJluN6KBGn7EiULpUqVi5CqcfawV2V+3a4STOWt1MoFKFCjd/tc+2VmB169QJsLeObNyI6CYn+AXYXoa+HQFrBRbsjjt/7N9Pgd1uzEXquLUCk98ejerXR4/uVNtGGuhg2Xev3f33GwjHEKElYP+SAc+/8groz/j6a2PdUK1KFcTardZRkkVvjNWBgcJAwH6BzV24EKD7DB5snnj/3r0RdObkyTBkGRzzlWFE8wTsF5gsS1ZNLeJ55uxZ86AHPfssgn48YQIMW1/QMM82zBHtF5jQH6zWn5k2a5YcNG+0bNIEQSe9+y4MWfjWfGU8R5RNPZeo1Sa/VKv21k9OhlvZb9ZzFAsKRkhg8pp609RU9Jy8lRRsR3Zs2xYVGKm2QpY1SYOd8CUr3mWtXIkafqdWHZflx/8+ceIWerLw4/wZM25JiuDXCAlMelcuxuTyTJJCYsgW4w937IgqtW3VCkaTRo1gNG7YEIYsW19GzYqQ9aHklQJZm/6k2rlThCEr2q9XO7ivVw/x5F+SnK8Q8f8/ZS/Z7UqW/5/f7tQoCkweUsmpTAai3Z1tpnWyTFWuei4iR8xUIFRRioaqNmYqI5sIjx0xAhHfUDcezFTA7ihXrlxBA/eqp+3yKNLuhhfYuigKTEC8Nnw47KU//3zDUJfskoeGZwI79uxBWQrMM8P4Lig7Cc2aOhUtaaMWDN1/6FB8ty0Etd+hNhB9PHr77gn+SJ/BhIJMvFg6ezYOdn7ySRjRWcFPaMTKkDNYrBzGox8KLE+vNVAbwy3LyEBCj+eeg7F73748WfnFiQAF9i+hKN5FdBoYedJl+bd+akNKmcibJx+/5CMgvwuOqgcA+bLYf4ACc9vH8lT6o88/R5nXx4+HcVptl+zWV8TynVCbR8sG09EBQIF572u5PHs7PR1eZnzzDQxZEdq7d4tKrsnMRGseaN3aoma5agoF5gqTy0x/HjyInJOmTYMxU13L/XX8uEsncZqtauXKqPkT3bvDeKpnTxhd1HyUCD5xpsD0jmd56rpo2TJEyszKgiE7A8rMEr1V8eFdNuxsn5ICNw916ABDJnO1V6+WysMPHwHtKUqBBd+XB7KzUYnft2yBsVEZ8ha2/Bw9dPgw8sjdF5nnJYZcLsqaTfmnKcoLoMm1asHhPWpyY36jhXoDQCZJogg/HQlQYI6ImIEEvBNI9F6UJUmABJwIUGBOhJhOAj4IUGA+4LEoCTgRoMCcCDGdBHwQoMB8wGNREnAiQIE5EWI6CfggQIH5gMeiJOBEgAJzIsR0EvBBgALzAY9FScCJAAXmRIjpJOCDAAXmAx6LkoATAQrMiRDTScAHAQrMBzwWJQEnAhSYEyGmk4APAhSYD3gsSgJOBCgwJ0JMJwEfBCgwH/BYlAScCFBgToSYTgI+CPwD+Cx4nR09n5kAAAAASUVORK5CYII=",
      "text/plain": []
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# order of axes matters as usual - you can repeat each element (pixel) 3 times\n",
    "# by changing order in parenthesis\n",
    "repeat(ims[0], \"h w c -> h (w repeat) c\", repeat=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note: `repeat` operation covers functionality identical to `numpy.repeat`, `numpy.tile` and actually more than that."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Reduce ⇆ repeat\n",
    "\n",
    "reduce and repeat are like opposite of each other: first one reduces amount of elements, second one increases.\n",
    "\n",
    "In the following example each image is repeated first, then we reduce over new axis to get back original tensor. Notice that operation patterns are \"reverse\" of each other"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "repeated = repeat(ims, \"b h w c -> b h new_axis w c\", new_axis=2)\n",
    "reduced = reduce(repeated, \"b h new_axis w c -> b h w c\", \"min\")\n",
    "assert numpy.array_equal(ims, reduced)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fancy examples in random order\n",
    "\n",
    "(a.k.a. mad designer gallery)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 11468 characters]",
      "text/plain": []
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# interweaving pixels of different pictures\n",
    "# all letters are observable\n",
    "rearrange(ims, \"(b1 b2) h w c -> (h b1) (w b2) c \", b1=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 14236 characters]",
      "text/plain": []
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# interweaving along vertical for couples of images\n",
    "rearrange(ims, \"(b1 b2) h w c -> (h b1) (b2 w) c\", b1=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 6680 characters]",
      "text/plain": []
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# interweaving lines for couples of images\n",
    "# exercise: achieve the same result without einops in your favourite framework\n",
    "reduce(ims, \"(b1 b2) h w c -> h (b2 w) c\", \"max\", b1=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 6468 characters]",
      "text/plain": []
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# color can be also composed into dimension\n",
    "# ... while image is downsampled\n",
    "reduce(ims, \"b (h 2) (w 2) c -> (c h) (b w)\", \"mean\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAMAAAAAYCAAAAACJ2q2ZAAAEV0lEQVR4AdVWfUxTVxQ//bJUBlXp6uaGCgUJAksYW2KiBkl0E5Nln0adzcBhZtZkX/zhR5a4zLDEOb8WE7fFuWSbMzqNNRRNVmNAYWEZ4GfiprKw0mEFSgtSBEpLd+5r332Pdx+blyVdvH+8e87v9zv3nnPfvfc9TQAevIVqfTmrDA+uT4ZSzzFJ/7JOgGfO8oRwjD5FqZYjrhHzh9YbHBFJkPIUUDANE8qYn4SsOKbQ8JyBy4d8OY4sjtGTIOUqIAn5cE8hP5GROtclX8RsW2LPZsfZ0ZLAHOUsSZBzTe1dPQNhMGeXVj6hKgn96L7eG7MsXLF2hgrvrm+/4x8M6zLylq2zqvAA3ccbb/pHDWmPL1hkTxEVsjfQ4vgjgeqrt2hEgdivdSesPRtEaGL/8gXqmw6+SG3JcG7uSzjpNXYJFi15/PZNIirrv986LHpGDzmOQpMOcf0LYv4Q2fWZyE+pH95Eh5Liv6kS84d77+6XcBVreNtuFm1+n+YPxTR/oAX0bQwDFDs7u5srMHi3VznCsUAgcFcJKv0yr7/9YwTD3yoZuLUNIdtxb5erCI2aK4yAAGWe3o5Dj6Kx8wbDO2N4BTo9ft+Vo++tkVhawOEg4UsfMeTtKwOInJYkHNa0VO2sd57GgFYm6OAYgPHUilTT4hPTAca/YAQEMKbpzK8eQWOcXQGy6CNBk9Y4d+VHZI0TjR7iegSi6wW0A59XE/wUuqJLAL1MXCMiyzMJbF1eC9DECCjwbEkbQAt1RaPihwEYetO8dHFpnkbEsKcF+NDp/5kyPJ8HGhQ3cIEBt6OikZKE/AHmoOlX0HLXhgWw8+f+tLUBYKCuDua+/nYa1dMtZKKQYPyHXzb5+kiDpqPZH3dJcsSdrN1HQpEOkS441VJTPhONzp1LyCaJN/oGCm8CFDWoTy6KDboowJDo8fUldwAaRo0YNHQeHyWTR4+R7ZOnxtscjtjv7gMB8G4+IfL0DbyByPXKazFC+OvU7xvNPCRPhvARuEd0PK0SxXc/wOqDb/WhuUE1NhTB3fdhN3LlDO/cUY8za/KFK0i6pOgbWFp1GMDlmjFv+khXD3y1mhmAAM99CXCtIH/c17WvQlUwOVhmPwJwrDY/+tsoil57XlXZlG0z3g4ilfsKw5927ddmWg0jPnJaCyhNC4BPLXvHcJv2C9Rl9QKqz+D3YfBXlLTxFgB7U74GuN8mDL9+j9Cxj9BVAZv9nUHJxX7By9XjicOPfUJpuoVAu6W1uiDuGp4ie0WlWdz2VAJrCwtV2H+G9LvcLwnBplW1B8ilrtJmC/OmV11kj0Ds88qseIS+eHtzLg2W/QsRbMQbjBgtc5j6qR4iHX2amZkmCeCxon/6YxlZ0luXx5J/oZVHe7xhS7ZOjsvsgb8Go8ZZT06oXjFYilSaLFBu6nP/VSKXT7R1NttEgPGsVgaSALNZskVL2kIi8pD1D30Bii30Py7/mkUAOfzz/w3PLyd1G8omRQAAAABJRU5ErkJggg==",
      "text/plain": []
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# disproportionate resize\n",
    "reduce(ims, \"b (h 4) (w 3) c -> (h) (b w)\", \"mean\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 4872 characters]",
      "text/plain": []
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# spilt each image in two halves, compute mean of the two\n",
    "reduce(ims, \"b (h1 h2) w c -> h2 (b w)\", \"mean\", h1=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 13716 characters]",
      "text/plain": []
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# split in small patches and transpose each patch\n",
    "rearrange(ims, \"b (h1 h2) (w1 w2) c -> (h1 w2) (b w1 h2) c\", h2=8, w2=8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 14888 characters]",
      "text/plain": []
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# stop me someone!\n",
    "rearrange(ims, \"b (h1 h2 h3) (w1 w2 w3) c -> (h1 w2 h3) (b w1 h2 w3) c\", h2=2, w2=2, w3=2, h3=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 12132 characters]",
      "text/plain": []
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rearrange(ims, \"(b1 b2) (h1 h2) (w1 w2) c -> (h1 b1 h2) (w1 b2 w2) c\", h1=3, w1=3, b2=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 17136 characters]",
      "text/plain": []
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# patterns can be arbitrarily complicated\n",
    "reduce(ims, \"(b1 b2) (h1 h2 h3) (w1 w2 w3) c -> (h1 w1 h3) (b1 w2 h2 w3 b2) c\", \"mean\", h2=2, w1=2, w3=2, h3=2, b2=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 9808 characters]",
      "text/plain": []
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# subtract background in each image individually and normalize\n",
    "# pay attention to () - this is composition of 0 axis, a dummy axis with 1 element.\n",
    "im2 = reduce(ims, \"b h w c -> b () () c\", \"max\") - ims\n",
    "im2 /= reduce(im2, \"b h w c -> b () () c\", \"max\")\n",
    "rearrange(im2, \"b h w c -> h (b w) c\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAkAAAABgCAIAAAB6y1p+AAAKyElEQVR4Ae3dfWxX1R3HcVpsGTS2UDREghprJQIN3SiI1ofxIE3JDCuKmDQpPqJbQuIYpQHpim4I1YQ2ghCWlLo/GimKESrRSDQVExq3QaVQH8pqK5ZSlGYpTz7UlY7s78+99q7n3Ht+17d/fs75nYfX+dkvNzn53aTB7u4R/OctcHqid5uJlrbWNjlMzUs1Ml+7ca3MM6/NlLn10DaQ9Q3YneCM5e+P3dXbH31Gs/05EnoGgHyPL9m3lUYEEEAAAQQcFaCAOXowLAsBBBBAwF+AAubvQysCCCCAgKMCFDBHD4ZlIYAAAgj4C1DA/H1oRQABBBBwVOAqR9f1s1lW7dZauddXa16V+dTcqTJ/dMWjMidEAAEE4irAE1hcT5Z9IYAAAjEXoIDF/IDZHgIIIBBXAQpYXE+WfSGAAAIxF6CAxfyA2R4CCCAQVwEKWFxPln0hgAACMRfgFmLEB3zTLTcFWkHWLVmB+tMZAQQQiKsAT2BxPVn2hQACCMRcgAIW8wNmewgggEBcBShgcT1Z9oUAAgjEXIACFvMDZnsIIIBAXAUoYHE9WfaFAAIIxFwgiTcy+5+w9RcOD+r5e071yIbrrr9O5pGF1oEi25mRiXkjsz8jLxz29xkBkC8QT2C+PDQigAACCLgqQAFz9WRYFwIIIICArwAFzJeHRgQQQAABVwUoYK6eDOtCAAEEEPAVoID58tCIAAIIIOCqgPXfQqzZtUvufV1lpcy/+/57mf9+2TKZV65bJ/MkmUYXLp2/VE5+9sxZmXuFxU8Uy6blK5fLPFHC/ftfk0utqvqzzDs722Wempoq8+XL/yDz0lI9vuzsYFhXVSVXVfv88zK/PKivvT5SVib7P7JmjcxdC3cf2C2XtP6v62XeeapT5mlj0mS+snilzMufKJd5UpJbf4EGPc59w4Yauf4tW+plfu7cRZlPmaJ/07W+flOg/rKzT8gTmA8OTQgggAAC7gpQwNw9G1aGAAIIIOAjQAHzwaEJAQQQQMBdAQqYu2fDyhBAAAEEfAQoYD44NCGAAAIIuCtg7BbiwY8+krt8cvVqmXvdipGdr4Qvbt8um3KnTZN5cVGRzKMKO050yKm9fvNQdr4S9n7T69WU0Hlr61G5/ra2T2TuFfb3/yCbqqv/IvO5cwtlnpeXL/OowoMNDXLq6lWrZB403Lp2rfzIpOxsmd+7ZInMowpbTrTIqdu+bJO5V9h/vl82VeyokPnkGyfL/KGCh2QeVfjee/+QU1dU7JB50PDYsX/Jj/T3/yhzUyFPYKYkGQcBBBBAIFQBClio3EyGAAIIIGBKgAJmSpJxEEAAAQRCFaCAhcrNZAgggAACpgQoYKYkGQcBBBBAIFQBY7cQt+zcKReempIi86a33pL5uIwMmc9auFDmu/bulblrtxAPdx2W6zz5xUmZ52e7dQtOLjLC8ODBT+XsR4/+U+ZPP/2wzJuaPpC5a7cQX9u2Ta4zeeRImb9y6JDMUzx+K7Lktttk/9c9bv+6dgtRLt4n/HSP/v709Oo3oc//3Xw52s69+u+ea7cQjxzR+5WbuhI2NFTLpoKC22Xe2Kj/vmVnXy/7mwp5AjMlyTgIIIAAAqEKUMBC5WYyBBBAAAFTAhQwU5KMgwACCCAQqgAFLFRuJkMAAQQQMCVAATMlyTgIIIAAAqEKGLuF+PePP5YLzxw7VuYNBw7I3Cscm54um9ra22VOGG+BlJRUucH8/Dky9wp7e7/xanIq/+ywvuU1JS9PrjNn9myZe4Ve43ze3Oz1kYTOUz2+P/NmzZP7mpY1TebH24/L3LUwK2tSoCWVl2+X/QcGBmS+aNGvZZ6cbPcZye7ockuECCCAAAIIDF+AAjZ8Q0ZAAAEEEIhAgAIWATpTIoAAAggMX4ACNnxDRkAAAQQQiECAAhYBOlMigAACCAxfwNgtxHPnz8vVXLx0SebPbd4s86DhxAkTgn6E/jEWGDXqF4F2NzDwn0D9o+p86cIFOXXG+PEyDxqmjxsnP/Ktx7yyc4zDCeP135mTZ/RvmbpG8cAD8+SSHn/8tzKvrW2Q+eLFpTLPz58u83feeVnm6elpMg8a8gQWVIz+CCCAAAJOCFDAnDgGFoEAAgggEFSAAhZUjP4IIIAAAk4IUMCcOAYWgQACCCAQVIACFlSM/ggggAACTggYu4Xo9VuF12Zmyo2+UF4u86Dh6NGjg37Eqf4jPd6o67XIywOXvZrI/yeQFEuHsddcI/d19vRpmQcNvz51Sn7Ea17ZOcZhb1+v3F1GWobMXQu9/s7U1FTIpa5aVSLzjRtrZV5X97bMd+zYI/OysodlHjTkCSyoGP0RQAABBJwQoIA5cQwsAgEEEEAgqAAFLKgY/RFAAAEEnBCggDlxDCwCAQQQQCCoAAUsqBj9EUAAAQScEDB2C/GOmTPlhvbs3y/zL7u6ZL7gnntkfqKzU+YTPG5nyc4OhuPG69+g81rqocZDsun4Ef1m2L5/98n+dy+4W+aEbgrk3nmnXNj7b7wh87fr6mT+Y3+/zL9obZX53MWLZZ7o4YfNH8otNH/eLHOvNy8X3F4g+7sW1te/K5fU06NvV86Zo/+eL1kyX47jdQuxq+tr2d9UyBOYKUnGQQABBBAIVYACFio3kyGAAAIImBKggJmSZBwEEEAAgVAFKGChcjMZAggggIApAQqYKUnGQQABBBAIVeAqU7OtWbFCDtXwrr79UrZhg+wfNKx69ln5kVm5uTJ3LUy7Wr+ZdHredLnUlsMtMi+cWSjzyVMny7zxk0aZE7opsKy0VC7sgzfflPmfSvRv2cnOV8LkZP1v2RKPeb3GSZT8seceM7LUkt8EczYy6f8xSFPTMfmprVvrZW4qvOuuX5oaSo6jv7WyKyECCCCAAALuCFDA3DkLVoIAAgggEECAAhYAi64IIIAAAu4IUMDcOQtWggACCCAQQIACFgCLrggggAAC7ggYu4U4IydH7qpxj34j5wvbtsn+HV99JfMbJk6U+cwEuW0oF+8TVtdWy9b1f1wv84vnL8q8sEjfTpSdIwxvvVV/f+6770G5qjFj9O3NUaNSZX+vcXJyfiX7uxbmzJ4tl7R53z6Z/62yUuaDg4MyX7Z6tcxz8/Nlnujh/fPul1vo6O6QedGcIpkXFxbL3LXwqaf0fvv6LsiltrSckHl6uv7/rrh4oey/dOkCmZsKeQIzJck4CCCAAAKhClDAQuVmMgQQQAABUwIUMFOSjIMAAgggEKoABSxUbiZDAAEEEDAlQAEzJck4CCCAAAKhCiQNdneHOmGiTXZaX35MtG3YWy9AvrZn+P74+szQL0D2/YxqfOblZ1Q8YtMrm2Tevq9d5jdPulnmkYWmgCLbgN2JeQKz68voCCCAAAKWBChglmAZFgEEEEDArgAFzK4voyOAAAIIWBKggFmCZVgEEEAAAbsCFDC7voyOAAIIIGBJgAJmCZZhEUAAAQTsClDA7PoyOgIIIICAJQEKmCVYhkUAAQQQsCtAAbPry+gIIIAAApYEKGCWYBkWAQQQQMCuAAXMri+jI4AAAghYEjD2RmZL62NYBBBA4CcFcrL1G70fvFe/0TtttH6z8E9ORAenBHgCc+o4WAwCCCCAwFAFKGBDlaIfAggggIBTAhQwp46DxSCAAAIIDFWAAjZUKfohgAACCDglQAFz6jhYDAIIIIDAUAX+Cxyhjd1yPQF/AAAAAElFTkSuQmCC",
      "text/plain": []
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# pixelate: first downscale by averaging, then upscale back using the same pattern\n",
    "averaged = reduce(ims, \"b (h h2) (w w2) c -> b h w c\", \"mean\", h2=6, w2=8)\n",
    "repeat(averaged, \"b h w c -> (h h2) (b w w2) c\", h2=6, w2=8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 11752 characters]",
      "text/plain": []
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rearrange(ims, \"b h w c -> w (b h) c\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 7860 characters]",
      "text/plain": []
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# let's bring color dimension as part of horizontal axis\n",
    "# at the same time horizontal axis is downsampled by 2x\n",
    "reduce(ims, \"b (h h2) (w w2) c -> (h w2) (b w c)\", \"mean\", h2=3, w2=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Ok, numpy is fun, but how do I use einops with some other framework?\n",
    "\n",
    "If that's what you've done with `ims` being numpy array:\n",
    "```python\n",
    "rearrange(ims, 'b h w c -> w (b h) c')\n",
    "```\n",
    "That's how you adapt the code for other frameworks:\n",
    "\n",
    "```python\n",
    "# pytorch:\n",
    "rearrange(ims, 'b h w c -> w (b h) c')\n",
    "# tensorflow:\n",
    "rearrange(ims, 'b h w c -> w (b h) c')\n",
    "# gluon:\n",
    "rearrange(ims, 'b h w c -> w (b h) c')\n",
    "# cupy:\n",
    "rearrange(ims, 'b h w c -> w (b h) c')\n",
    "# jax:\n",
    "rearrange(ims, 'b h w c -> w (b h) c')\n",
    "# paddle:\n",
    "rearrange(ims, 'b h w c -> w (b h) c')\n",
    "\n",
    "...well, you got the idea.\n",
    "```\n",
    "\n",
    "Einops allows backpropagation as if all operations were native to framework.\n",
    "Operations do not change when moving to another framework - einops notation is universal"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Summary\n",
    "\n",
    "- `rearrange` doesn't change number of elements and covers different numpy functions (like `transpose`, `reshape`, `stack`, `concatenate`,  `squeeze` and `expand_dims`)\n",
    "- `reduce` combines same reordering syntax with reductions (`mean`, `min`, `max`, `sum`, `prod`, and any others)\n",
    "- `repeat` additionally covers repeating and tiling\n",
    "- composition and decomposition of axes are a corner stone, they can and should be used together\n",
    "\n",
    "\n",
    "\n",
    "- [Second part of tutorial](https://github.com/arogozhnikov/einops/tree/main/docs) shows how einops works with other frameworks\n",
    "- [Third part of tutorial](https://arogozhnikov.github.io/einops/pytorch-examples.html) shows how to improve your DL code with einops"
   ]
  }
 ],
 "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.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
