{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Reinterpreting Tensors\n",
    "\n",
    "Sometimes the data in tensors needs to be interpreted as if it had different type or shape. For example, reading a binary file into memory produces a flat tensor of byte-valued data, which the application code may want to interpret as an array of data of specific shape and possibly different type.\n",
    "\n",
    "DALI provides the following operations which affect tensor metadata (shape, type, layout):\n",
    "* reshape\n",
    "* reinterpret\n",
    "* squeeze\n",
    "* expand_dims\n",
    "\n",
    "Thsese operations neither modify nor copy the data - the output tensor is just another view of the same region of memory, making these operations very cheap.\n",
    "\n",
    "## Fixed Output Shape\n",
    "\n",
    "This example demonstrates the simplest use of the `reshape` operation, assigning a new fixed shape to an existing tensor.\n",
    "\n",
    "First, we'll import DALI and other necessary modules, and define a utility for displaying the data, which will be used throughout this tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nvidia.dali as dali\n",
    "import nvidia.dali.fn as fn\n",
    "from nvidia.dali import pipeline_def\n",
    "import nvidia.dali.types as types\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "def show_result(outputs, names=[\"Input\", \"Output\"], formatter=None):\n",
    "    if not isinstance(outputs, tuple):\n",
    "        return show_result((outputs,))\n",
    "\n",
    "    outputs = [out.as_cpu() if hasattr(out, \"as_cpu\") else out for out in outputs]\n",
    "\n",
    "    for i in range(len(outputs[0])):\n",
    "        print(f\"---------------- Sample #{i} ----------------\")\n",
    "        for o, out in enumerate(outputs):\n",
    "            a = np.array(out[i])\n",
    "            s = \"x\".join(str(x) for x in a.shape)\n",
    "            title = names[o] if names is not None and o < len(names) else f\"Output #{o}\"\n",
    "            l = out.layout()\n",
    "            if l:\n",
    "                l += \" \"\n",
    "            print(f\"{title} ({l}{s})\")\n",
    "            np.set_printoptions(formatter=formatter)\n",
    "            print(a)\n",
    "\n",
    "\n",
    "def rand_shape(dims, lo, hi):\n",
    "    return list(np.random.randint(lo, hi, [dims]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's define out pipeline - it takes data from an external source and returns it both in original form and reshaped to a fixed square shape `[5, 5]`. Additionally, output tensors' layout is set to HW"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---------------- Sample #0 ----------------\n",
      "Input (25)\n",
      "[3 6 5 4 8 9 1 7 9 6 8 0 5 0 9 6 2 0 5 2 6 3 7 0 9]\n",
      "Output (HW 5x5)\n",
      "[[3 6 5 4 8]\n",
      " [9 1 7 9 6]\n",
      " [8 0 5 0 9]\n",
      " [6 2 0 5 2]\n",
      " [6 3 7 0 9]]\n",
      "---------------- Sample #1 ----------------\n",
      "Input (25)\n",
      "[0 3 2 3 1 3 1 3 7 1 7 4 0 5 1 5 9 9 4 0 9 8 8 6 8]\n",
      "Output (HW 5x5)\n",
      "[[0 3 2 3 1]\n",
      " [3 1 3 7 1]\n",
      " [7 4 0 5 1]\n",
      " [5 9 9 4 0]\n",
      " [9 8 8 6 8]]\n",
      "---------------- Sample #2 ----------------\n",
      "Input (25)\n",
      "[6 3 1 2 5 2 5 6 7 4 3 5 6 4 6 2 4 2 7 9 7 7 2 9 7]\n",
      "Output (HW 5x5)\n",
      "[[6 3 1 2 5]\n",
      " [2 5 6 7 4]\n",
      " [3 5 6 4 6]\n",
      " [2 4 2 7 9]\n",
      " [7 7 2 9 7]]\n"
     ]
    }
   ],
   "source": [
    "@pipeline_def(device_id=0, num_threads=4, batch_size=3)\n",
    "def example1(input_data):\n",
    "    np.random.seed(1234)\n",
    "    inp = fn.external_source(input_data, batch=False, dtype=types.INT32)\n",
    "    return inp, fn.reshape(inp, shape=[5, 5], layout=\"HW\")\n",
    "\n",
    "\n",
    "pipe1 = example1(lambda: np.random.randint(0, 10, size=[25], dtype=np.int32))\n",
    "pipe1.build()\n",
    "show_result(pipe1.run())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see, the numbers from flat input tensors have been rearranged into 5x5 matrices.\n",
    "\n",
    "## Reshape with Wildcards\n",
    "\n",
    "Let's now consider a more advanced use case. Imagine you have some flattened array that represents a fixed number of columns, but the number of rows is free to vary from sample to sample. In that case, you can put a wildcard dimension by specifying its shape as `-1`. Whe using wildcards, the output is resized so that the total number of elements is the same as in the input."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---------------- Sample #0 ----------------\n",
      "Input (25)\n",
      "[5 1 4 9 5 2 1 6 1 9 7 6 0 2 9 1 2 6 7 7 7 8 7 1 7]\n",
      "Output (5x5)\n",
      "[[5 1 4 9 5]\n",
      " [2 1 6 1 9]\n",
      " [7 6 0 2 9]\n",
      " [1 2 6 7 7]\n",
      " [7 8 7 1 7]]\n",
      "---------------- Sample #1 ----------------\n",
      "Input (35)\n",
      "[0 3 5 7 3 1 5 2 5 3 8 5 2 5 3 0 6 8 0 5 6 8 9 2 2 2 9 7 5 7 1 0 9 3 0]\n",
      "Output (7x5)\n",
      "[[0 3 5 7 3]\n",
      " [1 5 2 5 3]\n",
      " [8 5 2 5 3]\n",
      " [0 6 8 0 5]\n",
      " [6 8 9 2 2]\n",
      " [2 9 7 5 7]\n",
      " [1 0 9 3 0]]\n",
      "---------------- Sample #2 ----------------\n",
      "Input (30)\n",
      "[0 6 2 1 5 8 6 5 1 0 5 8 2 9 4 7 9 5 2 4 8 2 5 6 5 9 6 1 9 5]\n",
      "Output (6x5)\n",
      "[[0 6 2 1 5]\n",
      " [8 6 5 1 0]\n",
      " [5 8 2 9 4]\n",
      " [7 9 5 2 4]\n",
      " [8 2 5 6 5]\n",
      " [9 6 1 9 5]]\n"
     ]
    }
   ],
   "source": [
    "@pipeline_def(device_id=0, num_threads=4, batch_size=3)\n",
    "def example2(input_data):\n",
    "    np.random.seed(12345)\n",
    "    inp = fn.external_source(input_data, batch=False, dtype=types.INT32)\n",
    "    return inp, fn.reshape(inp, shape=[-1, 5])\n",
    "\n",
    "\n",
    "pipe2 = example2(\n",
    "    lambda: np.random.randint(0, 10, size=[5 * np.random.randint(3, 10)], dtype=np.int32)\n",
    ")\n",
    "pipe2.build()\n",
    "show_result(pipe2.run())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Removing and Adding Unit Dimensions\n",
    "\n",
    "There are two dedicated operators `squeeze` and `expand_dims` which can be used for removing and adding dimensions with unit extent. The following example demonstrates the removal of a redundant dimension as well as adding two new dimensions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---------------- Sample #0 ----------------\n",
      "Input (CHW 1x6x3)\n",
      "[[[8 2 1]\n",
      "  [7 5 9]\n",
      "  [2 4 6]\n",
      "  [0 8 6]\n",
      "  [5 3 1]\n",
      "  [1 6 1]]]\n",
      "Output (HW 6x3)\n",
      "[[8 2 1]\n",
      " [7 5 9]\n",
      " [2 4 6]\n",
      " [0 8 6]\n",
      " [5 3 1]\n",
      " [1 6 1]]\n",
      "Output #2 (FHWC 1x6x3x1)\n",
      "[[[[8]\n",
      "   [2]\n",
      "   [1]]\n",
      "\n",
      "  [[7]\n",
      "   [5]\n",
      "   [9]]\n",
      "\n",
      "  [[2]\n",
      "   [4]\n",
      "   [6]]\n",
      "\n",
      "  [[0]\n",
      "   [8]\n",
      "   [6]]\n",
      "\n",
      "  [[5]\n",
      "   [3]\n",
      "   [1]]\n",
      "\n",
      "  [[1]\n",
      "   [6]\n",
      "   [1]]]]\n",
      "---------------- Sample #1 ----------------\n",
      "Input (CHW 1x2x2)\n",
      "[[[6 9]\n",
      "  [0 9]]]\n",
      "Output (HW 2x2)\n",
      "[[6 9]\n",
      " [0 9]]\n",
      "Output #2 (FHWC 1x2x2x1)\n",
      "[[[[6]\n",
      "   [9]]\n",
      "\n",
      "  [[0]\n",
      "   [9]]]]\n",
      "---------------- Sample #2 ----------------\n",
      "Input (CHW 1x2x6)\n",
      "[[[4 4 6 6 6 3]\n",
      "  [8 2 1 7 9 7]]]\n",
      "Output (HW 2x6)\n",
      "[[4 4 6 6 6 3]\n",
      " [8 2 1 7 9 7]]\n",
      "Output #2 (FHWC 1x2x6x1)\n",
      "[[[[4]\n",
      "   [4]\n",
      "   [6]\n",
      "   [6]\n",
      "   [6]\n",
      "   [3]]\n",
      "\n",
      "  [[8]\n",
      "   [2]\n",
      "   [1]\n",
      "   [7]\n",
      "   [9]\n",
      "   [7]]]]\n"
     ]
    }
   ],
   "source": [
    "@pipeline_def(device_id=0, num_threads=4, batch_size=3)\n",
    "def example_squeeze_expand(input_data):\n",
    "    np.random.seed(4321)\n",
    "    inp = fn.external_source(input_data, batch=False, layout=\"CHW\", dtype=types.INT32)\n",
    "    squeezed = fn.squeeze(inp, axes=[0])\n",
    "    expanded = fn.expand_dims(squeezed, axes=[0, 3], new_axis_names=\"FC\")\n",
    "    return inp, fn.squeeze(inp, axes=[0]), expanded\n",
    "\n",
    "\n",
    "def single_channel_generator():\n",
    "    return np.random.randint(0, 10, size=[1] + rand_shape(2, 1, 7), dtype=np.int32)\n",
    "\n",
    "\n",
    "pipe_squeeze_expand = example_squeeze_expand(single_channel_generator)\n",
    "pipe_squeeze_expand.build()\n",
    "show_result(pipe_squeeze_expand.run())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Rearranging Dimensions\n",
    "\n",
    "Reshape allows you to swap, insert or remove dimenions. The argument `src_dims` allows you to specify which source dimension is used for a given output dimension. You can also insert a new dimension by specifying -1 as a source dimension index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---------------- Sample #0 ----------------\n",
      "Input (6x3)\n",
      "[[8 2 1]\n",
      " [7 5 9]\n",
      " [2 4 6]\n",
      " [0 8 6]\n",
      " [5 3 1]\n",
      " [1 6 1]]\n",
      "Output (3x6)\n",
      "[[8 2 1 7 5 9]\n",
      " [2 4 6 0 8 6]\n",
      " [5 3 1 1 6 1]]\n",
      "---------------- Sample #1 ----------------\n",
      "Input (2x2)\n",
      "[[6 9]\n",
      " [0 9]]\n",
      "Output (2x2)\n",
      "[[6 9]\n",
      " [0 9]]\n",
      "---------------- Sample #2 ----------------\n",
      "Input (2x6)\n",
      "[[4 4 6 6 6 3]\n",
      " [8 2 1 7 9 7]]\n",
      "Output (6x2)\n",
      "[[4 4]\n",
      " [6 6]\n",
      " [6 3]\n",
      " [8 2]\n",
      " [1 7]\n",
      " [9 7]]\n"
     ]
    }
   ],
   "source": [
    "@pipeline_def(device_id=0, num_threads=4, batch_size=3)\n",
    "def example_reorder(input_data):\n",
    "    np.random.seed(4321)\n",
    "    inp = fn.external_source(input_data, batch=False, dtype=types.INT32)\n",
    "    return inp, fn.reshape(inp, src_dims=[1, 0])\n",
    "\n",
    "\n",
    "pipe_reorder = example_reorder(\n",
    "    lambda: np.random.randint(0, 10, size=rand_shape(2, 1, 7), dtype=np.int32)\n",
    ")\n",
    "pipe_reorder.build()\n",
    "show_result(pipe_reorder.run())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Adding and Removing Dimensions\n",
    "\n",
    "Dimensions can be added or removed by specifying `src_dims` argument or by using dedicated `squeeze` and `expand_dims` operators.\n",
    "\n",
    "The following example reinterprets single-channel data from CHW to HWC layout by discarding the leading dimension and adding a new trailing dimension. It also specifies the output layout."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---------------- Sample #0 ----------------\n",
      "Input (CHW 1x4x3)\n",
      "[[[2 8 2]\n",
      "  [1 7 5]\n",
      "  [9 2 4]\n",
      "  [6 0 8]]]\n",
      "Output (HWC 4x3x1)\n",
      "[[[2]\n",
      "  [8]\n",
      "  [2]]\n",
      "\n",
      " [[1]\n",
      "  [7]\n",
      "  [5]]\n",
      "\n",
      " [[9]\n",
      "  [2]\n",
      "  [4]]\n",
      "\n",
      " [[6]\n",
      "  [0]\n",
      "  [8]]]\n",
      "---------------- Sample #1 ----------------\n",
      "Input (CHW 1x4x3)\n",
      "[[[6 5 3]\n",
      "  [1 1 6]\n",
      "  [1 1 9]\n",
      "  [6 9 0]]]\n",
      "Output (HWC 4x3x1)\n",
      "[[[6]\n",
      "  [5]\n",
      "  [3]]\n",
      "\n",
      " [[1]\n",
      "  [1]\n",
      "  [6]]\n",
      "\n",
      " [[1]\n",
      "  [1]\n",
      "  [9]]\n",
      "\n",
      " [[6]\n",
      "  [9]\n",
      "  [0]]]\n",
      "---------------- Sample #2 ----------------\n",
      "Input (CHW 1x4x3)\n",
      "[[[9 9 5]\n",
      "  [4 4 6]\n",
      "  [6 6 3]\n",
      "  [8 2 1]]]\n",
      "Output (HWC 4x3x1)\n",
      "[[[9]\n",
      "  [9]\n",
      "  [5]]\n",
      "\n",
      " [[4]\n",
      "  [4]\n",
      "  [6]]\n",
      "\n",
      " [[6]\n",
      "  [6]\n",
      "  [3]]\n",
      "\n",
      " [[8]\n",
      "  [2]\n",
      "  [1]]]\n"
     ]
    }
   ],
   "source": [
    "@pipeline_def(device_id=0, num_threads=4, batch_size=3)\n",
    "def example_remove_add(input_data):\n",
    "    np.random.seed(4321)\n",
    "    inp = fn.external_source(input_data, batch=False, layout=\"CHW\", dtype=types.INT32)\n",
    "    return inp, fn.reshape(\n",
    "        inp, src_dims=[1, 2, -1], layout=\"HWC\"  # select HW and add a new one at the end\n",
    "    )  # specify the layout string\n",
    "\n",
    "\n",
    "pipe_remove_add = example_remove_add(lambda: np.random.randint(0, 10, [1, 4, 3], dtype=np.int32))\n",
    "pipe_remove_add.build()\n",
    "show_result(pipe_remove_add.run())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Relative Shape\n",
    "\n",
    "The output shape may be calculated in relative terms, with a new extent being a multiple of a source extent.\n",
    "For example, you may want to combine two subsequent rows into one - doubling the number of columns and halving the number of rows. The use of relative shape can be combined with dimension rearranging, in which case the new output extent is a multiple of a _different_ source extent.\n",
    "\n",
    "The example below reinterprets the input as having twice as many _columns_ as the input had _rows_."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---------------- Sample #0 ----------------\n",
      "Input (4x6)\n",
      "[[5 4 8 9 1 7]\n",
      " [9 6 8 0 5 0]\n",
      " [9 6 2 0 5 2]\n",
      " [6 3 7 0 9 0]]\n",
      "Output (3x8)\n",
      "[[5 4 8 9 1 7 9 6]\n",
      " [8 0 5 0 9 6 2 0]\n",
      " [5 2 6 3 7 0 9 0]]\n",
      "---------------- Sample #1 ----------------\n",
      "Input (4x6)\n",
      "[[3 1 3 1 3 7]\n",
      " [1 7 4 0 5 1]\n",
      " [5 9 9 4 0 9]\n",
      " [8 8 6 8 6 3]]\n",
      "Output (3x8)\n",
      "[[3 1 3 1 3 7 1 7]\n",
      " [4 0 5 1 5 9 9 4]\n",
      " [0 9 8 8 6 8 6 3]]\n",
      "---------------- Sample #2 ----------------\n",
      "Input (2x6)\n",
      "[[5 2 5 6 7 4]\n",
      " [3 5 6 4 6 2]]\n",
      "Output (3x4)\n",
      "[[5 2 5 6]\n",
      " [7 4 3 5]\n",
      " [6 4 6 2]]\n"
     ]
    }
   ],
   "source": [
    "@pipeline_def(device_id=0, num_threads=4, batch_size=3)\n",
    "def example_rel_shape(input_data):\n",
    "    np.random.seed(1234)\n",
    "    inp = fn.external_source(input_data, batch=False, dtype=types.INT32)\n",
    "    return inp, fn.reshape(inp, rel_shape=[0.5, 2], src_dims=[1, 0])\n",
    "\n",
    "\n",
    "pipe_rel_shape = example_rel_shape(\n",
    "    lambda: np.random.randint(\n",
    "        0, 10, [np.random.randint(1, 7), 2 * np.random.randint(1, 5)], dtype=np.int32\n",
    "    )\n",
    ")\n",
    "\n",
    "pipe_rel_shape.build()\n",
    "show_result(pipe_rel_shape.run())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reinterpreting Data Type\n",
    "\n",
    "The `reinterpret` operation can view the data as if it was of different type. When a new shape is not specified, the innermost dimension is resized accordingly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---------------- Sample #0 ----------------\n",
      "Input (4x12)\n",
      "[[0x35 0xdc 0x5d 0xd1 0xcc 0xec 0x0e 0x70 0x74 0x5d 0xb3 0x9c]\n",
      " [0x98 0x42 0x0d 0xc9 0xf9 0xd7 0x77 0xc5 0x8f 0x7e 0xac 0xc7]\n",
      " [0xb1 0xda 0x54 0xdc 0x17 0xa1 0xc8 0x45 0xe9 0x24 0x90 0x26]\n",
      " [0x9a 0x5c 0xc6 0x46 0x1e 0x20 0xd2 0x32 0xab 0x7e 0x47 0xcd]]\n",
      "Output (4x3)\n",
      "[[0xd15ddc35 0x700eeccc 0x9cb35d74]\n",
      " [0xc90d4298 0xc577d7f9 0xc7ac7e8f]\n",
      " [0xdc54dab1 0x45c8a117 0x269024e9]\n",
      " [0x46c65c9a 0x32d2201e 0xcd477eab]]\n",
      "---------------- Sample #1 ----------------\n",
      "Input (5x4)\n",
      "[[0x1a 0x1f 0x3d 0xe0]\n",
      " [0x76 0x35 0xbb 0x1d]\n",
      " [0xba 0xe9 0x99 0x5b]\n",
      " [0x78 0xe8 0x4d 0x03]\n",
      " [0x70 0x37 0x41 0x80]]\n",
      "Output (5x1)\n",
      "[[0xe03d1f1a]\n",
      " [0x1dbb3576]\n",
      " [0x5b99e9ba]\n",
      " [0x034de878]\n",
      " [0x80413770]]\n",
      "---------------- Sample #2 ----------------\n",
      "Input (5x8)\n",
      "[[0x50 0x6d 0xbd 0x54 0xc9 0xa3 0x73 0xb6]\n",
      " [0x7f 0xc9 0x79 0xcd 0xf6 0xc0 0xc8 0x5e]\n",
      " [0xfe 0x09 0x27 0x19 0xaf 0x8d 0xaa 0x8f]\n",
      " [0x32 0x96 0x55 0x0e 0xf0 0x0e 0xca 0x80]\n",
      " [0xfb 0x56 0x52 0x71 0x4c 0x54 0x86 0x03]]\n",
      "Output (5x2)\n",
      "[[0x54bd6d50 0xb673a3c9]\n",
      " [0xcd79c97f 0x5ec8c0f6]\n",
      " [0x192709fe 0x8faa8daf]\n",
      " [0x0e559632 0x80ca0ef0]\n",
      " [0x715256fb 0x0386544c]]\n"
     ]
    }
   ],
   "source": [
    "@pipeline_def(device_id=0, num_threads=4, batch_size=3)\n",
    "def example_reinterpret(input_data):\n",
    "    np.random.seed(1234)\n",
    "    inp = fn.external_source(input_data, batch=False, dtype=types.UINT8)\n",
    "    return inp, fn.reinterpret(inp, dtype=dali.types.UINT32)\n",
    "\n",
    "\n",
    "pipe_reinterpret = example_reinterpret(\n",
    "    lambda: np.random.randint(\n",
    "        0, 255, [np.random.randint(1, 7), 4 * np.random.randint(1, 5)], dtype=np.uint8\n",
    "    )\n",
    ")\n",
    "\n",
    "pipe_reinterpret.build()\n",
    "\n",
    "\n",
    "def hex_bytes(x):\n",
    "    f = f\"0x{{:0{2*x.nbytes}x}}\"\n",
    "    return f.format(x)\n",
    "\n",
    "\n",
    "show_result(pipe_reinterpret.run(), formatter={\"int\": hex_bytes})"
   ]
  }
 ],
 "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}