tutorials/indexing.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Indexing, Views, and Accessing Values\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Copyright (c) Meta Platforms, Inc. and affiliates.**\n",
    "\n",
    "This source code is licensed under the MIT license found in the\n",
    "LICENSE file in the root directory of this source tree."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Introduction\n",
    "\n",
    "This notebook will discuss accessing data in a tensor through indexing, views, and accessing the actual `Float` values at a location in a tensor. \n",
    "\n",
    "Indexing is when one takes an individual value, as a `DTensor`, from the tensor. Indexing can also be used to take a section of a tensor and the section can be a 1D vector, a 2D array, or higher dimensional tensor. Both indexing and taking a section are done using the index operator, `[]`. An alternative way to index a tensor is to use the `view()` function.\n",
    "\n",
    "The values returned in indexing or views are of type DTensor or DScalar. Converting the values to a float will also be discussed."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Housekeeping\n",
    "\n",
    "This notebook uses `api.jar` from the diffkt project.<br>\n",
    "`@file:DependsOn(\"...\")` tells the Kotlin Jupyter notebook the path to a jar that it needs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "@file:DependsOn(\"../kotlin/api/build/libs/api.jar\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import org.diffkt.*\n",
    "import java.util.Arrays"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Indexing a Tensor\n",
    "\n",
    "Indexing for tensors starts at 0. The example we will use is a 2x3x4 tensor. Indexing an element or taking a section of the tensor uses the index operator, `[]`. Let the indexes for the individual axis be `i1`, `i2`, and `i3`, such that an index for `x` is `x[i1,i2, i3]`. Below is an example of an index that returns an individual element of a tensor."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Example 3D Tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0], [9.0, 10.0, 11.0, 12.0]], [[13.0, 14.0, 15.0, 16.0], [17.0, 18.0, 19.0, 20.0], [21.0, 22.0, 23.0, 24.0]]]\n"
     ]
    }
   ],
   "source": [
    "// example 3D tensor\n",
    "\n",
    "val x = tensorOf(1f, 2f, 3f, 4f, \n",
    "                 5f, 6f, 7f, 8f, \n",
    "                 9f, 10f, 11f, 12f, \n",
    "                 13f, 14f, 15f, 16f, \n",
    "                 17f, 18f, 19f, 20f, \n",
    "                 21f, 22f, 23f, 24f).reshape(2,3,4)\n",
    "\n",
    "println(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Example of Retrieving Individual Elements"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Indexing along i3\n",
      "x[0,0,0] = 1.0\n",
      "x[0,0,1] = 2.0\n",
      "x[0,0,2] = 3.0\n",
      "x[0,0,3] = 4.0\n",
      "\n",
      "Indexing along i2\n",
      "x[0,0,0] = 1.0\n",
      "x[0,1,0] = 5.0\n",
      "x[0,2,0] = 9.0\n",
      "\n",
      "Indexing along i1\n",
      "x[0,0,0] = 1.0\n",
      "x[1,0,0] = 13.0\n"
     ]
    }
   ],
   "source": [
    "// indexing along the different axes for individual elements of 3D tensor\n",
    "\n",
    "println(\"Indexing along i3\")\n",
    "println(\"x[0,0,0] = ${x[0,0,0]}\")\n",
    "println(\"x[0,0,1] = ${x[0,0,1]}\")\n",
    "println(\"x[0,0,2] = ${x[0,0,2]}\")\n",
    "println(\"x[0,0,3] = ${x[0,0,3]}\")\n",
    "println(\"\")\n",
    "\n",
    "println(\"Indexing along i2\")\n",
    "println(\"x[0,0,0] = ${x[0,0,0]}\")\n",
    "println(\"x[0,1,0] = ${x[0,1,0]}\")\n",
    "println(\"x[0,2,0] = ${x[0,2,0]}\")\n",
    "println(\"\")\n",
    "\n",
    "println(\"Indexing along i1\")\n",
    "println(\"x[0,0,0] = ${x[0,0,0]}\")\n",
    "println(\"x[1,0,0] = ${x[1,0,0]}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Taking a 1D section\n",
    "\n",
    "When we take a section of a 3D tensor with a 2D index, the indexing operator returns a 1D vector. In this example a section is taken with `i1` and `i2` as the index. A 1D vector is returned with the values of the tensor along the `i3` axis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x[0,0] = [1.0, 2.0, 3.0, 4.0]\n",
      "x[0,1] = [5.0, 6.0, 7.0, 8.0]\n",
      "x[0,2] = [9.0, 10.0, 11.0, 12.0]\n",
      "x[1,0] = [13.0, 14.0, 15.0, 16.0]\n",
      "x[1,1] = [17.0, 18.0, 19.0, 20.0]\n",
      "x[1,2] = [21.0, 22.0, 23.0, 24.0]\n"
     ]
    }
   ],
   "source": [
    "// indexing to produce a 1D section\n",
    "\n",
    "println(\"x[0,0] = ${x[0,0]}\")\n",
    "println(\"x[0,1] = ${x[0,1]}\")\n",
    "println(\"x[0,2] = ${x[0,2]}\")\n",
    "println(\"x[1,0] = ${x[1,0]}\")\n",
    "println(\"x[1,1] = ${x[1,1]}\")\n",
    "println(\"x[1,2] = ${x[1,2]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Taking a 2D Section\n",
    "\n",
    "When we take a section with a 1D index of a 3D tensor, the indexing operator returns a 2D array from the tensor. In this example `i1` is fixed and a 2D array is returned with the values indexed by `i2` and `i3`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x[0] = [[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0], [9.0, 10.0, 11.0, 12.0]]\n",
      "x[1] = [[13.0, 14.0, 15.0, 16.0], [17.0, 18.0, 19.0, 20.0], [21.0, 22.0, 23.0, 24.0]]\n"
     ]
    }
   ],
   "source": [
    "// indexing to produce a 2D section\n",
    "\n",
    "println(\"x[0] = ${x[0]}\")\n",
    "println(\"x[1] = ${x[1]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Indexing a Transposed Tensor\n",
    "\n",
    "If you want to reverse the order of the sectioning, start with the transpose of the tensor. This reverses the order of the axes. Let the indexes of the transposed tensor be `t1`, `t2`, and `t3`, such that `t1 = i3`, `t2 = t2`, and `t3 = i1`. Note, the shape of the transposed tensor is (4,3,2)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Example of Transposed 3D Tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[[1.0, 13.0], [5.0, 17.0], [9.0, 21.0]], [[2.0, 14.0], [6.0, 18.0], [10.0, 22.0]], [[3.0, 15.0], [7.0, 19.0], [11.0, 23.0]], [[4.0, 16.0], [8.0, 20.0], [12.0, 24.0]]]\n"
     ]
    }
   ],
   "source": [
    "// transpose of example 3D tensor\n",
    "\n",
    "val x = tensorOf(1f, 2f, 3f, 4f, \n",
    "                 5f, 6f, 7f, 8f, \n",
    "                 9f, 10f, 11f, 12f, \n",
    "                 13f, 14f, 15f, 16f, \n",
    "                 17f, 18f, 19f, 20f, \n",
    "                 21f, 22f, 23f, 24f).reshape(2,3,4).transpose()\n",
    "\n",
    "println(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Example of Retrieving Individual Elements of a Transposed Tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Indexing along t3 or i1\n",
      "x[0,0,0] = 1.0\n",
      "x[0,0,1] = 13.0\n",
      "\n",
      "Indexing along t2 or i2\n",
      "x[0,0,0] = 1.0\n",
      "x[0,1,0] = 5.0\n",
      "x[0,2,0] = 9.0\n",
      "\n",
      "Indexing along t1 or i3\n",
      "x[0,0,0] = 1.0\n",
      "x[1,0,0] = 2.0\n",
      "x[2,0,0] = 3.0\n",
      "x[3,0,0] = 4.0\n"
     ]
    }
   ],
   "source": [
    "// indexing along the different axes for individual elements of transposed 3D tensor\n",
    "\n",
    "println(\"Indexing along t3 or i1\")\n",
    "println(\"x[0,0,0] = ${x[0,0,0]}\")\n",
    "println(\"x[0,0,1] = ${x[0,0,1]}\")\n",
    "println(\"\")\n",
    "\n",
    "println(\"Indexing along t2 or i2\")\n",
    "println(\"x[0,0,0] = ${x[0,0,0]}\")\n",
    "println(\"x[0,1,0] = ${x[0,1,0]}\")\n",
    "println(\"x[0,2,0] = ${x[0,2,0]}\")\n",
    "println(\"\")\n",
    "\n",
    "println(\"Indexing along t1 or i3\")\n",
    "println(\"x[0,0,0] = ${x[0,0,0]}\")\n",
    "println(\"x[1,0,0] = ${x[1,0,0]}\")\n",
    "println(\"x[2,0,0] = ${x[2,0,0]}\")\n",
    "println(\"x[3,0,0] = ${x[3,0,0]}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Taking a 1D Section of a Transposed Tensor\n",
    "\n",
    "When we take a section with a 2D index of a 3D tensor, the indexing operator returns a 1D vector. In this example a section is taken with `t1` and `t2` as the index, which is the equivalent of `i3` and `i2`. A 1D vector is returned of the values of the tensor along the `t3` or `i1` axis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x[0,0] = [1.0, 13.0]\n",
      "x[0,1] = [5.0, 17.0]\n",
      "x[0,2] = [9.0, 21.0]\n",
      "x[1,0] = [2.0, 14.0]\n",
      "x[1,1] = [6.0, 18.0]\n",
      "x[1,2] = [10.0, 22.0]\n",
      "x[2,0] = [3.0, 15.0]\n",
      "x[2,1] = [7.0, 19.0]\n",
      "x[2,2] = [11.0, 23.0]\n",
      "x[3,0] = [4.0, 16.0]\n",
      "x[3,1] = [8.0, 20.0]\n",
      "x[3,2] = [12.0, 24.0]\n"
     ]
    }
   ],
   "source": [
    "// indexing to produce a 1D section\n",
    "\n",
    "println(\"x[0,0] = ${x[0,0]}\")\n",
    "println(\"x[0,1] = ${x[0,1]}\")\n",
    "println(\"x[0,2] = ${x[0,2]}\")\n",
    "println(\"x[1,0] = ${x[1,0]}\")\n",
    "println(\"x[1,1] = ${x[1,1]}\")\n",
    "println(\"x[1,2] = ${x[1,2]}\")\n",
    "println(\"x[2,0] = ${x[2,0]}\")\n",
    "println(\"x[2,1] = ${x[2,1]}\")\n",
    "println(\"x[2,2] = ${x[2,2]}\")\n",
    "println(\"x[3,0] = ${x[3,0]}\")\n",
    "println(\"x[3,1] = ${x[3,1]}\")\n",
    "println(\"x[3,2] = ${x[3,2]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Taking a 2D Section of a Transposed Tensor\n",
    "\n",
    "When we take a section with a 1D index of a 3D tensor, the indexing operator returns a 2D array from the tensor. In this example where t1, or i3, is fixed, a 2D array is returned with the values indexed by t2 and t3, or i2 and i1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x[0] = [[1.0, 13.0], [5.0, 17.0], [9.0, 21.0]]\n",
      "x[1] = [[2.0, 14.0], [6.0, 18.0], [10.0, 22.0]]\n",
      "x[2] = [[3.0, 15.0], [7.0, 19.0], [11.0, 23.0]]\n",
      "x[3] = [[4.0, 16.0], [8.0, 20.0], [12.0, 24.0]]\n"
     ]
    }
   ],
   "source": [
    "// indexing to produce a 2D section\n",
    "\n",
    "println(\"x[0] = ${x[0]}\")\n",
    "println(\"x[1] = ${x[1]}\")\n",
    "println(\"x[2] = ${x[2]}\")\n",
    "println(\"x[3] = ${x[3]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Using Views\n",
    "\n",
    "In addition to using the indexing operator, `[]`, we can call the `view()` function to retrieve data from a tensor."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Example 3D Tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0], [9.0, 10.0, 11.0, 12.0]], [[13.0, 14.0, 15.0, 16.0], [17.0, 18.0, 19.0, 20.0], [21.0, 22.0, 23.0, 24.0]]]\n"
     ]
    }
   ],
   "source": [
    "// example 3D tensor\n",
    "\n",
    "val x = tensorOf(1f, 2f, 3f, 4f, \n",
    "                 5f, 6f, 7f, 8f, \n",
    "                 9f, 10f, 11f, 12f, \n",
    "                 13f, 14f, 15f, 16f, \n",
    "                 17f, 18f, 19f, 20f, \n",
    "                 21f, 22f, 23f, 24f).reshape(2,3,4)\n",
    "\n",
    "println(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Using `view()` with Indices\n",
    "\n",
    "The `view()` function is called with an IntArray of indices, in a similar manner as the the index operator, `[]`.\n",
    "\n",
    "The first example gets an individual value from the tensor using an `intArray` of indices the same length as the dimension of the tensor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "view along i3\n",
      "x.view(intArrayOf(0,0,0)) = 1.0\n",
      "x.view(intArrayOf(0,0,1)) = 2.0\n",
      "x.view(intArrayOf(0,0,2)) = 3.0\n",
      "x.view(intArrayOf(0,0,3)) = 4.0\n",
      "\n",
      "view along i2\n",
      "x.view(intArrayOf(0,0,0)) = 1.0\n",
      "x.view(intArrayOf(0,1,0)) = 5.0\n",
      "x.view(intArrayOf(0,2,0)) = 9.0\n",
      "\n",
      "view along i1\n",
      "x.view(intArrayOf(0,0,0)) = 1.0\n",
      "x.view(intArrayOf(1,0,0)) = 13.0\n"
     ]
    }
   ],
   "source": [
    "// view of individual elements of 3D tensor\n",
    "\n",
    "println(\"view along i3\")\n",
    "println(\"x.view(intArrayOf(0,0,0)) = ${x.view(intArrayOf(0,0,0))}\")\n",
    "println(\"x.view(intArrayOf(0,0,1)) = ${x.view(intArrayOf(0,0,1))}\")\n",
    "println(\"x.view(intArrayOf(0,0,2)) = ${x.view(intArrayOf(0,0,2))}\")\n",
    "println(\"x.view(intArrayOf(0,0,3)) = ${x.view(intArrayOf(0,0,3))}\")\n",
    "println(\"\")\n",
    "\n",
    "println(\"view along i2\")\n",
    "println(\"x.view(intArrayOf(0,0,0)) = ${x.view(intArrayOf(0,0,0))}\")\n",
    "println(\"x.view(intArrayOf(0,1,0)) = ${x.view(intArrayOf(0,1,0))}\")\n",
    "println(\"x.view(intArrayOf(0,2,0)) = ${x.view(intArrayOf(0,2,0))}\")\n",
    "println(\"\")\n",
    "\n",
    "println(\"view along i1\")\n",
    "println(\"x.view(intArrayOf(0,0,0)) = ${x.view(intArrayOf(0,0,0))}\")\n",
    "println(\"x.view(intArrayOf(1,0,0)) = ${x.view(intArrayOf(1,0,0))}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Taking a 1D Section\n",
    "\n",
    "Using 2 indices returns a 1D vector section, in a similar manner as the index operator, `[]`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x.view(intArrayOf(0,0)) = [1.0, 2.0, 3.0, 4.0]\n",
      "x.view(intArrayOf(0,1)) = [5.0, 6.0, 7.0, 8.0]\n",
      "x.view(intArrayOf(0,2)) = [9.0, 10.0, 11.0, 12.0]\n",
      "x.view(intArrayOf(1,0)) = [13.0, 14.0, 15.0, 16.0]\n",
      "x.view(intArrayOf(1,1)) = [17.0, 18.0, 19.0, 20.0]\n",
      "x.view(intArrayOf(1,2)) = [21.0, 22.0, 23.0, 24.0]\n"
     ]
    }
   ],
   "source": [
    "// get a 1D section\n",
    "\n",
    "println(\"x.view(intArrayOf(0,0)) = ${x.view(intArrayOf(0,0))}\")\n",
    "println(\"x.view(intArrayOf(0,1)) = ${x.view(intArrayOf(0,1))}\")\n",
    "println(\"x.view(intArrayOf(0,2)) = ${x.view(intArrayOf(0,2))}\")\n",
    "println(\"x.view(intArrayOf(1,0)) = ${x.view(intArrayOf(1,0))}\")\n",
    "println(\"x.view(intArrayOf(1,1)) = ${x.view(intArrayOf(1,1))}\")\n",
    "println(\"x.view(intArrayOf(1,2)) = ${x.view(intArrayOf(1,2))}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Taking a 2D Section\n",
    "\n",
    "Using one indice returns a 2D array section, in a similar manner to the index operator, `[]`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x.view(intArrayOf(0)) = [[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0], [9.0, 10.0, 11.0, 12.0]]\n",
      "x.view(intArrayOf(1)) = [[13.0, 14.0, 15.0, 16.0], [17.0, 18.0, 19.0, 20.0], [21.0, 22.0, 23.0, 24.0]]\n"
     ]
    }
   ],
   "source": [
    "// get a 2D section\n",
    "\n",
    "println(\"x.view(intArrayOf(0)) = ${x.view(intArrayOf(0))}\")\n",
    "println(\"x.view(intArrayOf(1)) = ${x.view(intArrayOf(1))}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "In all three cases, the `view()` function performed exactly the same as the index operator, `[]`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Using `view()` with an Axis\n",
    "\n",
    "The signature for this version of `view()` is:\n",
    "\n",
    "`fun DTensor.view(index: Int, axis: Int): DTensor`\n",
    "\n",
    "For a 3D tensor, this function returns a 2D array section such that the axis selected is fixed for the value of the index. Remember we start indexing from 0. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "axis 0\n",
      "x.view(0,0) = [[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0], [9.0, 10.0, 11.0, 12.0]]\n",
      "x.view(1,0) = [[13.0, 14.0, 15.0, 16.0], [17.0, 18.0, 19.0, 20.0], [21.0, 22.0, 23.0, 24.0]]\n",
      "\n",
      "axis 1\n",
      "x.view(0,1) = [[1.0, 2.0, 3.0, 4.0], [13.0, 14.0, 15.0, 16.0]]\n",
      "x.view(1,1) = [[5.0, 6.0, 7.0, 8.0], [17.0, 18.0, 19.0, 20.0]]\n",
      "x.view(2,1) = [[9.0, 10.0, 11.0, 12.0], [21.0, 22.0, 23.0, 24.0]]\n",
      "\n",
      "axis 2\n",
      "x.view(0,2) = [[1.0, 5.0, 9.0], [13.0, 17.0, 21.0]]\n",
      "x.view(1,2) = [[2.0, 6.0, 10.0], [14.0, 18.0, 22.0]]\n",
      "x.view(2,2) = [[3.0, 7.0, 11.0], [15.0, 19.0, 23.0]]\n",
      "x.view(3,2) = [[4.0, 8.0, 12.0], [16.0, 20.0, 24.0]]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "// 2D section along an axis\n",
    "\n",
    "println(\"axis 0\")\n",
    "println(\"x.view(0,0) = ${x.view(0,0)}\")\n",
    "println(\"x.view(1,0) = ${x.view(1,0)}\")\n",
    "println(\"\")\n",
    "\n",
    "println(\"axis 1\")\n",
    "println(\"x.view(0,1) = ${x.view(0,1)}\")\n",
    "println(\"x.view(1,1) = ${x.view(1,1)}\")\n",
    "println(\"x.view(2,1) = ${x.view(2,1)}\")\n",
    "\n",
    "\n",
    "println(\"\")\n",
    "\n",
    "println(\"axis 2\")\n",
    "println(\"x.view(0,2) = ${x.view(0,2)}\")\n",
    "println(\"x.view(1,2) = ${x.view(1,2)}\")\n",
    "println(\"x.view(2,2) = ${x.view(2,2)}\")\n",
    "println(\"x.view(3,2) = ${x.view(3,2)}\")\n",
    "println(\"\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Using `view()` with an Index Being a Range\n",
    "\n",
    "This example is better explained with a larger 3x3x3 tensor.\n",
    "\n",
    "The signature for this version of `view()` is:\n",
    "\n",
    "`fun DTensor.view(index: IntRange, axis: Int): DTensor`\n",
    "\n",
    "This implementation of `view()` returns a 2D array section, but it can return one or more section along the axis that is fixed by passing a range for the index for the sections. Because the function can return more than one section, the returned tensor has the same rank, or number of dimensions, as the tensor being sectioned. \n",
    "\n",
    "For example, when the range of the index is a single number, it is equivant the call to `view(index: Int, axis: Int)` with a single index,\n",
    "\n",
    "`x.view(0..0,0)` = `x.view(0,0)`\n",
    "\n",
    "except the returned tensor has Shape(1,3,3) instead of Shape(3,3).\n",
    "\n",
    "When the range has two numbers, 0..1, 1..2, then two sections are taken and returned in a tensor of Shape(2,3,3), Shape(3,2,3), or Shape(3,3,2), depending on the axis that was fixed.\n",
    "\n",
    "When the range equals the number of dimensions or rank of the tensor, 0..2, then the complete tensor of Shape(3,3,3) is returned."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Example 3x3x3 Tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "// a larger tensor\n",
    "\n",
    "val x = tensorOf(1f, 2f, 3f, 4f, 5f, 6f, 7f, 8f, 9f, \n",
    "                 10f, 11f, 12f, 13f, 14f, 15f, 16f, 17f, \n",
    "                 18f, 19f, 20f, 21f, 22f, 23f, 24f, 25f, 26f, 27f).reshape(Shape(3,3,3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Take One 2D Section\n",
    "\n",
    "Take one 2D section along the axis that is indicated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x.view(0..0,0) = [[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]]\n",
      "x.view(1..1,0) = [[[10.0, 11.0, 12.0], [13.0, 14.0, 15.0], [16.0, 17.0, 18.0]]]\n",
      "x.view(2..2,0) = [[[19.0, 20.0, 21.0], [22.0, 23.0, 24.0], [25.0, 26.0, 27.0]]]\n",
      "x.view(0..0,1) = [[[1.0, 2.0, 3.0]], [[10.0, 11.0, 12.0]], [[19.0, 20.0, 21.0]]]\n",
      "x.view(1..1,1) = [[[4.0, 5.0, 6.0]], [[13.0, 14.0, 15.0]], [[22.0, 23.0, 24.0]]]\n",
      "x.view(2..2,1) = [[[7.0, 8.0, 9.0]], [[16.0, 17.0, 18.0]], [[25.0, 26.0, 27.0]]]\n",
      "x.view(0..0,2) = [[[1.0], [4.0], [7.0]], [[10.0], [13.0], [16.0]], [[19.0], [22.0], [25.0]]]\n",
      "x.view(1..1,2) = [[[2.0], [5.0], [8.0]], [[11.0], [14.0], [17.0]], [[20.0], [23.0], [26.0]]]\n",
      "x.view(2..2,2) = [[[3.0], [6.0], [9.0]], [[12.0], [15.0], [18.0]], [[21.0], [24.0], [27.0]]]\n"
     ]
    }
   ],
   "source": [
    "// take one 2D section\n",
    "\n",
    "println(\"x.view(0..0,0) = ${x.view(0..0,0)}\")\n",
    "println(\"x.view(1..1,0) = ${x.view(1..1,0)}\")\n",
    "println(\"x.view(2..2,0) = ${x.view(2..2,0)}\")\n",
    "println(\"x.view(0..0,1) = ${x.view(0..0,1)}\")\n",
    "println(\"x.view(1..1,1) = ${x.view(1..1,1)}\")\n",
    "println(\"x.view(2..2,1) = ${x.view(2..2,1)}\")\n",
    "println(\"x.view(0..0,2) = ${x.view(0..0,2)}\")\n",
    "println(\"x.view(1..1,2) = ${x.view(1..1,2)}\")\n",
    "println(\"x.view(2..2,2) = ${x.view(2..2,2)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Take Two 2D Sections\n",
    "\n",
    "Take two 2D sections along the axis indicated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x.view(0..1,0) = [[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]], [[10.0, 11.0, 12.0], [13.0, 14.0, 15.0], [16.0, 17.0, 18.0]]]\n",
      "x.view(1..2,0) = [[[10.0, 11.0, 12.0], [13.0, 14.0, 15.0], [16.0, 17.0, 18.0]], [[19.0, 20.0, 21.0], [22.0, 23.0, 24.0], [25.0, 26.0, 27.0]]]\n",
      "x.view(0..1,1) = [[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], [[10.0, 11.0, 12.0], [13.0, 14.0, 15.0]], [[19.0, 20.0, 21.0], [22.0, 23.0, 24.0]]]\n",
      "x.view(1..2,1) = [[[4.0, 5.0, 6.0], [7.0, 8.0, 9.0]], [[13.0, 14.0, 15.0], [16.0, 17.0, 18.0]], [[22.0, 23.0, 24.0], [25.0, 26.0, 27.0]]]\n",
      "x.view(0..1,2) = [[[1.0, 2.0], [4.0, 5.0], [7.0, 8.0]], [[10.0, 11.0], [13.0, 14.0], [16.0, 17.0]], [[19.0, 20.0], [22.0, 23.0], [25.0, 26.0]]]\n",
      "x.view(1..2,2) = [[[2.0, 3.0], [5.0, 6.0], [8.0, 9.0]], [[11.0, 12.0], [14.0, 15.0], [17.0, 18.0]], [[20.0, 21.0], [23.0, 24.0], [26.0, 27.0]]]\n"
     ]
    }
   ],
   "source": [
    "// take two 2D sections\n",
    "\n",
    "println(\"x.view(0..1,0) = ${x.view(0..1,0)}\")\n",
    "println(\"x.view(1..2,0) = ${x.view(1..2,0)}\")\n",
    "println(\"x.view(0..1,1) = ${x.view(0..1,1)}\")\n",
    "println(\"x.view(1..2,1) = ${x.view(1..2,1)}\")\n",
    "println(\"x.view(0..1,2) = ${x.view(0..1,2)}\")\n",
    "println(\"x.view(1..2,2) = ${x.view(1..2,2)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Take Three 2D Sections\n",
    "\n",
    "Take three 2D sections along the axis indicated. Regardless of the axis selected, the complete tensor is returned."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x.view(0..2,0) = [[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]], [[10.0, 11.0, 12.0], [13.0, 14.0, 15.0], [16.0, 17.0, 18.0]], [[19.0, 20.0, 21.0], [22.0, 23.0, 24.0], [25.0, 26.0, 27.0]]]\n",
      "x.view(0..2,1) = [[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]], [[10.0, 11.0, 12.0], [13.0, 14.0, 15.0], [16.0, 17.0, 18.0]], [[19.0, 20.0, 21.0], [22.0, 23.0, 24.0], [25.0, 26.0, 27.0]]]\n",
      "x.view(0..2,2) = [[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]], [[10.0, 11.0, 12.0], [13.0, 14.0, 15.0], [16.0, 17.0, 18.0]], [[19.0, 20.0, 21.0], [22.0, 23.0, 24.0], [25.0, 26.0, 27.0]]]\n"
     ]
    }
   ],
   "source": [
    "// take three 2D sections\n",
    "\n",
    "println(\"x.view(0..2,0) = ${x.view(0..2,0)}\")\n",
    "println(\"x.view(0..2,1) = ${x.view(0..2,1)}\")\n",
    "println(\"x.view(0..2,2) = ${x.view(0..2,2)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Accessing the internal value of a `DScalar` or `DTensor`\n",
    "\n",
    "When you use indexing, a `DScalar` or `DTensor` is returned. \n",
    "\n",
    "To access the internal value of a `DScalar`, which is a `Float`, one needs to call the `value` property."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "y = 1.0, type = Float\n"
     ]
    }
   ],
   "source": [
    "// Create a FloatScalar and access its value\n",
    "\n",
    "val x = FloatScalar(1f)\n",
    "val y = x.value \n",
    "\n",
    "println(\"y = ${y}, type = ${y::class.simpleName}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "To access the internal value of a `DTensor` at a particular index, use the `at(index)` function on the tensor. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "y = 1.0, type = Float\n"
     ]
    }
   ],
   "source": [
    "// Access the internal value of a tensor at an index position\n",
    "\n",
    "val x = tensorOf(1f, 2f, 3f)\n",
    "val y = x.at(0)\n",
    "\n",
    "println(\"y = ${y}, type = ${y::class.simpleName}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "To access the internal values of a `DTensor` for every value of the tensor or a subset of the tensor, one has to write a function for the conversion. \n",
    "\n",
    "Below is an example of converting a 1D tensor to a `FloatArray`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "y = [1.0, 2.0, 3.0], type = FloatArray\n"
     ]
    }
   ],
   "source": [
    "// Convert a 1D tensor to a FloatArray\n",
    "\n",
    "val x = tensorOf(1f, 2f, 3f)\n",
    "val range = 0..(x.size - 1)\n",
    "val y = range.map{index -> x.at(index)}.toFloatArray()\n",
    "\n",
    "println(\"y = ${Arrays.toString(y)}, type = ${y::class.simpleName}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# The End"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Kotlin",
   "language": "kotlin",
   "name": "kotlin"
  },
  "language_info": {
   "codemirror_mode": "text/x-kotlin",
   "file_extension": ".kt",
   "mimetype": "text/x-kotlin",
   "name": "kotlin",
   "nbconvert_exporter": "",
   "pygments_lexer": "kotlin",
   "version": "1.7.0-dev-3303"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}