{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Installing packages:\n",
      "\t.package(path: \"/home/ubuntu/fastai_docs/dev_swift/FastaiNotebook_00_load_data\")\n",
      "\t\tFastaiNotebook_00_load_data\n",
      "With SwiftPM flags: []\n",
      "Working in: /tmp/tmpxb804zjr\n",
      "Fetching https://github.com/mxcl/Path.swift\n",
      "Fetching https://github.com/JustHTTP/Just\n",
      "Completed resolution in 1.20s\n",
      "Cloning https://github.com/mxcl/Path.swift\n",
      "Resolving https://github.com/mxcl/Path.swift at 0.16.2\n",
      "Cloning https://github.com/JustHTTP/Just\n",
      "Resolving https://github.com/JustHTTP/Just at 0.7.1\n",
      "Compile Swift Module 'Path' (9 sources)\n",
      "Compile Swift Module 'Just' (1 sources)\n",
      "Compile Swift Module 'FastaiNotebook_00_load_data' (1 sources)\n",
      "Compile Swift Module 'jupyterInstalledPackages' (1 sources)\n",
      "Linking ./.build/x86_64-unknown-linux/debug/libjupyterInstalledPackages.so\n",
      "Initializing Swift...\n",
      "Loading library...\n",
      "Installation complete!\n"
     ]
    }
   ],
   "source": [
    "%install '.package(path: \"$cwd/FastaiNotebook_00_load_data\")' FastaiNotebook_00_load_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import FastaiNotebook_00_load_data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import Path\n",
    "import TensorFlow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "let (xTrain, yTrain, xValid, yValid) = loadMNIST(path: mnistPath, flat: true)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import Python"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "let np = Python.import(\"numpy\")\n",
    "let plt = Python.import(\"matplotlib.pyplot\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('inline', 'module://ipykernel.pylab.backend_inline')\n"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%include \"EnableIPythonDisplay.swift\"\n",
    "IPythonDisplay.shell.enable_matplotlib(\"inline\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "let img = np.array(xTrain[0].array.scalars).reshape(28,28)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAATUAAAEyCAYAAACbGke8AAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDMuMC4yLCBodHRwOi8vbWF0cGxvdGxpYi5vcmcvOIA7rQAAD61JREFUeJzt3X+o1XWex/HXa63+yCyV2UycWqcIW4v2tpgtFVsRTj+YqFvNMkKDS5H9kWAwyIb/TP1hyFbOIkXokI3FjNNA02SxbEVaLrRIV7My3bYIp9EuSplp9gu97/3jfoNr4/X78Zxz7znnfZ8PkHvO9778nPfpW6++58f3HEeEACCLv2n3AADQSpQagFQoNQCpUGoAUqHUAKRCqQFIhVIDkAqlBiAVSg1AKieM5o3Z5vQFAI36JCL+ti7EkRqAbvHnklBTpWb7Wtvv2f7A9r3NrAUArdBwqdkeJ+lRSddJmilpru2ZrRoMABrRzJHabEkfRMSHEfGtpN9LurE1YwFAY5optWmS/jLk+s5q2xFsz7fdZ7uvidsCgCLNvPrpo2z7q1c3I2KlpJUSr34CGHnNHKntlHTmkOs/lPRxc+MAQHOaKbU3JJ1r+0e2T5L0M0lrWzMWADSm4YefEXHI9gJJL0oaJ2lVRLzbsskAoAEeze8o4Dk1AE3YFBGz6kKcUQAgFUoNQCqUGoBUKDUAqVBqAFKh1ACkQqkBSIVSA5AKpQYgFUoNQCqUGoBUKDUAqVBqAFKh1ACkQqkBSIVSA5AKpQYgFUoNQCqUGoBUKDUAqVBqAFKh1ACkQqkBSIVSA5AKpQYgFUoNQCqUGoBUKDUAqVBqAFKh1ACkQqkBSIVSA5AKpQYgFUoNQCqUGoBUKDUAqVBqAFKh1ACkckK7B0B3GzduXG3mtNNOG4VJjrRgwYKi3Mknn1yUmzFjRlHu7rvvrs089NBDRWvNnTu3KPf111/XZpYuXVq01v3331+U62RNlZrtHZIOSDos6VBEzGrFUADQqFYcqV0VEZ+0YB0AaBrPqQFIpdlSC0kv2d5ke/7RArbn2+6z3dfkbQFArWYffl4WER/bPl3Sy7b/NyI2DA1ExEpJKyXJdjR5ewBwTE0dqUXEx9XPPZKelTS7FUMBQKMaLjXb421P+O6ypB9L2tqqwQCgEc08/Jwi6Vnb363zu4j4r5ZMBQANarjUIuJDSf/QwlkwjLPOOqs2c9JJJxWtdemllxblLr/88qLcxIkTazO33HJL0VqdbOfOnUW55cuX12Z6e3uL1jpw4EBR7q233qrNvPbaa0VrZcBbOgCkQqkBSIVSA5AKpQYgFUoNQCqUGoBUKDUAqVBqAFKh1ACk4ojR++AMPqXjSD09PUW5devW1Wba8ZHZGQwMDBTlbr/99qLcF1980cw4R+jv7y/KffbZZ7WZ9957r9lxOsGmkk/X5kgNQCqUGoBUKDUAqVBqAFKh1ACkQqkBSIVSA5AKpQYgFUoNQCrNfu8nmvDRRx8V5T799NPaTIYzCjZu3FiU27dvX23mqquuKlrr22+/Lco99dRTRTm0H0dqAFKh1ACkQqkBSIVSA5AKpQYgFUoNQCqUGoBUKDUAqfDm2zbau3dvUW7RokW1mZ/85CdFa7355ptFueXLlxflSmzZsqUoN2fOnKLcwYMHazPnn39+0VoLFy4syqF7cKQGIBVKDUAqlBqAVCg1AKlQagBSodQApEKpAUiFUgOQCqUGIBVHxOjdmD16NzbGnHrqqUW5AwcOFOVWrFhRlLvjjjtqM7fddlvRWmvWrCnKYczaFBGz6kIcqQFIpbbUbK+yvcf21iHbJtt+2fb71c9JIzsmAJQpOVL7jaRrv7ftXkmvRMS5kl6prgNA29WWWkRskPT9j5O4UdLq6vJqSTe1eC4AaEijHz00JSL6JSki+m2fPlzQ9nxJ8xu8HQA4LiP+eWoRsVLSSolXPwGMvEZf/dxte6okVT/3tG4kAGhco6W2VtK86vI8Sc+1ZhwAaE7JWzrWSPofSTNs77R9h6SlkubYfl/SnOo6ALRd7XNqETF3mF9d3eJZ0IT9+/e3dL3PP/+8ZWvdeeedRbmnn366KDcwMNDMOEiOMwoApEKpAUiFUgOQCqUGIBVKDUAqlBqAVCg1AKlQagBSodQApMJ3FOCoxo8fX5R7/vnnazNXXHFF0VrXXXddUe6ll14qyiEdvqMAwNhDqQFIhVIDkAqlBiAVSg1AKpQagFQoNQCpUGoAUuHNt2jKOeecU5vZvHlz0Vr79u0ryq1fv74209fXV7TWo48+WpQbzf9OMCzefAtg7KHUAKRCqQFIhVIDkAqlBiAVSg1AKpQagFQoNQCpUGoAUuGMAoy43t7eotwTTzxRlJswYUIz4xxh8eLFRbknn3yyKNff39/MODg2zigAMPZQagBSodQApEKpAUiFUgOQCqUGIBVKDUAqlBqAVCg1AKlwRgE6xgUXXFCUW7ZsWW3m6quvbnacI6xYsaIot2TJktrMrl27mh1nrGrNGQW2V9neY3vrkG332d5le0v15/pmpwWAVih5+PkbSdceZfuvIqKn+vOfrR0LABpTW2oRsUHS3lGYBQCa1swLBQtsv109PJ00XMj2fNt9tsu+iBEAmtBoqT0m6RxJPZL6JT08XDAiVkbErJIn+ACgWQ2VWkTsjojDETEg6deSZrd2LABoTEOlZnvqkKu9krYOlwWA0XRCXcD2GklXSvqB7Z2SfinpSts9kkLSDkl3jeCMAFCMN9+i60ycOLE2c8MNNxStVfoR4raLcuvWravNzJkzp2gt/BU+zhvA2EOpAUiFUgOQCqUGIBVKDUAqlBqAVCg1AKlQagBSodQApMIZBRjTvvnmm6LcCSfUnlEoSTp06FBt5pprrila69VXXy3KjSGcUQBg7KHUAKRCqQFIhVIDkAqlBiAVSg1AKpQagFQoNQCpUGoAUil7mzQwCi688MKi3K233lqbufjii4vWKj1ToNS2bdtqMxs2bGjpbeJIHKkBSIVSA5AKpQYgFUoNQCqUGoBUKDUAqVBqAFKh1ACkQqkBSIUzCtCUGTNm1GYWLFhQtNbNN99clDvjjDOKcq10+PDholx/f39tZmBgoNlxcAwcqQFIhVIDkAqlBiAVSg1AKpQagFQoNQCpUGoAUqHUAKTCm2/HmNI3rs6dO7coV/LG2unTpxet1Q59fX1FuSVLlhTl1q5d28w4aAGO1ACkUltqts+0vd72dtvv2l5YbZ9s+2Xb71c/J438uABwbCVHaock/SIi/l7SP0m62/ZMSfdKeiUizpX0SnUdANqqttQioj8iNleXD0jaLmmapBslra5iqyXdNFJDAkCp43qhwPZ0SRdJ2ihpSkT0S4PFZ/v0Yf7OfEnzmxsTAMoUl5rtUyQ9I+meiNhvu+jvRcRKSSurNaKRIQGgVNGrn7ZP1GCh/TYi/lht3m17avX7qZL2jMyIAFCu5NVPS3pc0vaIWDbkV2slzasuz5P0XOvHA4DjU/Lw8zJJP5f0ju0t1bbFkpZK+oPtOyR9JOmnIzMiAJRzxOg9zcVzao2ZMmVKbWbmzJlFaz3yyCNFufPOO68o1w4bN26szTz44INFaz33XNkDDD6CuyNsiohZdSHOKACQCqUGIBVKDUAqlBqAVCg1AKlQagBSodQApEKpAUiFUgOQCt9RMAImT55clFuxYkVRrqenpzZz9tlnF63VDq+//npR7uGHHy7Kvfjii7WZr776qmgt5MORGoBUKDUAqVBqAFKh1ACkQqkBSIVSA5AKpQYgFUoNQCq8+bZyySWXFOUWLVpUm5k9e3bRWtOmTSvKtcOXX35ZlFu+fHlt5oEHHiha6+DBg0U54Fg4UgOQCqUGIBVKDUAqlBqAVCg1AKlQagBSodQApEKpAUiFUgOQCmcUVHp7e1uaa6Vt27bVZl544YWitQ4dOlSUK/1o7X379hXlgNHCkRqAVCg1AKlQagBSodQApEKpAUiFUgOQCqUGIBVKDUAqlBqAVBwRo3dj9ujdGIBsNkXErLpQ7ZGa7TNtr7e93fa7thdW2++zvcv2lurP9a2YGgCaUXLu5yFJv4iIzbYnSNpk++Xqd7+KiIdGbjwAOD61pRYR/ZL6q8sHbG+X1Lnf7QZgTDuuFwpsT5d0kaSN1aYFtt+2vcr2pBbPBgDHrbjUbJ8i6RlJ90TEfkmPSTpHUo8Gj+SO+lk1tufb7rPd14J5AeCYil79tH2ipBckvRgRy47y++mSXoiIC2rW4dVPAI1q2auflvS4pO1DC8321CGxXklbG5kSAFqp5NXPyyT9XNI7trdU2xZLmmu7R1JI2iHprhGZEACOA2++BdAtWvPwEwC6CaUGIBVKDUAqlBqAVCg1AKlQagBSodQApEKpAUiFUgOQCqUGIBVKDUAqlBqAVCg1AKlQagBSodQApEKpAUiFUgOQCqUGIBVKDUAqJV+80kqfSPrz97b9oNrerbp9fqn770O3zy91/30Yjfn/riQ0ql+8ctQB7L6SL1PoVN0+v9T996Hb55e6/z500vw8/ASQCqUGIJVOKLWV7R6gSd0+v9T996Hb55e6/z50zPxtf04NAFqpE47UAKBlKDUAqbSt1Gxfa/s92x/YvrddczTD9g7b79jeYruv3fOUsL3K9h7bW4dsm2z7ZdvvVz8ntXPGYxlm/vts76r2wxbb17dzxmOxfabt9ba3237X9sJqezftg+HuQ0fsh7Y8p2Z7nKT/kzRH0k5Jb0iaGxHbRn2YJtjeIWlWRHTNmyZt/7OkLyQ9GREXVNv+XdLeiFha/Q9mUkT8WzvnHM4w898n6YuIeKids5WwPVXS1IjYbHuCpE2SbpL0r+qefTDcffgXdcB+aNeR2mxJH0TEhxHxraTfS7qxTbOMKRGxQdLe722+UdLq6vJqDf4L2pGGmb9rRER/RGyuLh+QtF3SNHXXPhjuPnSEdpXaNEl/GXJ9pzroH8pxCEkv2d5ke367h2nClIjolwb/hZV0epvnacQC229XD0879qHbULanS7pI0kZ16T743n2QOmA/tKvUfJRt3fjekssi4h8lXSfp7uqhEUbfY5LOkdQjqV/Sw+0dp57tUyQ9I+meiNjf7nkacZT70BH7oV2ltlPSmUOu/1DSx22apWER8XH1c4+kZzX4sLob7a6eJ/nu+ZI9bZ7nuETE7og4HBEDkn6tDt8Ptk/UYBn8NiL+WG3uqn1wtPvQKfuhXaX2hqRzbf/I9kmSfiZpbZtmaYjt8dWTpLI9XtKPJW099t/qWGslzasuz5P0XBtnOW7flUGlVx28H2xb0uOStkfEsiG/6pp9MNx96JT90LYzCqqXe/9D0jhJqyJiSVsGaZDtszV4dCYNfoTT77rhPtheI+lKDX5UzG5Jv5T0J0l/kHSWpI8k/TQiOvLJ+GHmv1KDD3lC0g5Jd333/FSnsX25pP+W9I6kgWrzYg0+J9Ut+2C4+zBXHbAfOE0KQCqcUQAgFUoNQCqUGoBUKDUAqVBqAFKh1ACkQqkBSOX/AUiz/VGRMZv/AAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 360x360 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "None\n"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "plt.figure(figsize: [5,5])\n",
    "let fig = plt.imshow(X:img, cmap:\"gray\")\n",
    "plt.show(fig)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Matmul"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "var weights = Tensor<Float>(randomNormal: [784, 10]) / sqrt(784)\n",
    "var bias = Tensor(repeating: 0.0, shape: [10])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### From scratch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Flattened version in pure swift"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "func swiftMatmul(a:[Float], b:[Float], aShape:[Int], bShape:[Int]) -> [Float]{\n",
    "    var res:[Float] = Array(repeating: 0.0, count: aShape[0] * bShape[1])\n",
    "    for i in 0..<aShape[0]{\n",
    "        for j in 0..<bShape[1]{\n",
    "            for k in 0..<aShape[1]{\n",
    "                res[i + aShape[0] * j] += a[i + aShape[0]*k] * b[k + bShape[0]*j]\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "    return res\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "let flatA = xTrain[0..<5].array.scalars\n",
    "let flatB = weights.array.scalars"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "let res = swiftMatmul(a: flatA, b:flatB, aShape:[5,784], bShape:[784,10])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.23496165999999996 ms\r\n"
     ]
    }
   ],
   "source": [
    "time(repeating: 100) { let _ = swiftMatmul(a: flatA, b:flatB, aShape:[5,784], bShape:[784,10])}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With pointers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "func swiftMatmultPointers(a:[Float], b:[Float], aShape:[Int], bShape:[Int]) -> [Float]{\n",
    "    var res:[Float] = Array(repeating: 0.0, count: aShape[0] * bShape[1])\n",
    "    res.withUnsafeMutableBufferPointer { res_ in\n",
    "        a.withUnsafeBufferPointer { a_ in\n",
    "            b.withUnsafeBufferPointer { b_ in                              \n",
    "                for i in 0..<aShape[0]{\n",
    "                    for j in 0..<bShape[1]{\n",
    "                        for k in 0..<aShape[1]{\n",
    "                            res_[i + aShape[0] * j] += a_[i + aShape[0]*k] * b_[k + bShape[0]*j]\n",
    "                        }\n",
    "                    }\n",
    "                }\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "    return res\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.14819732 ms\r\n"
     ]
    }
   ],
   "source": [
    "time(repeating: 100) { let _ = swiftMatmultPointers(a: flatA, b:flatB, aShape:[5,784], bShape:[784,10])}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With `ShapedArray`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "func shapedArrayMatmul(_ a:ShapedArray<Float>, _ b:ShapedArray<Float>) -> ShapedArray<Float>{\n",
    "    var res:ShapedArray<Float> = ShapedArray(repeating: 0.0, shape: [a.shape[0], b.shape[1]])\n",
    "    for i in 0..<a.shape[0]{\n",
    "        for j in 0..<b.shape[1]{\n",
    "            for k in 0..<a.shape[1]{\n",
    "                res[i][j].scalar! += a[i][k].scalar! * b[k][j].scalar!\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "    return res\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "let m1 = xTrain[0..<5].reshaped(toShape: [5, 784])\n",
    "let m2 = weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3680.167613 ms\r\n"
     ]
    }
   ],
   "source": [
    "time() { let res = shapedArrayMatmul(m1.array, m2.array)}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With `Tensor`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "func TensorMatmul(_ a:Tensor<Float>, _ b:Tensor<Float>) -> Tensor<Float>{\n",
    "    var res:Tensor<Float> = Tensor(repeating: 0.0, shape: [a.shape[0], b.shape[1]])\n",
    "    for i in 0..<a.shape[0]{\n",
    "        for j in 0..<b.shape[1]{\n",
    "            for k in 0..<a.shape[1]{\n",
    "                res[i][j] += a[i][k] * b[k][j]\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "    return res\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "error: Execution was interrupted, reason: signal SIGSTOP.\nThe process has been left at the point where it was interrupted, use \"thread return -x\" to return to the state before expression evaluation.\n"
     ]
    }
   ],
   "source": [
    "time() { let res = TensorMatmul(m1, m2)}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Too long to execture -> Looping over `ShapedArray` or `Tensor` indices is a bad idea!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Elementwise ops"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Operators (+,-,\\*,/) are usually element-wise.\n",
    "\n",
    "Examples of element-wise operations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "▿ 2 elements\n",
       "  - .0 : [10.0, 6.0, -4.0]\n",
       "  - .1 : [2.0, 8.0, 7.0]\n"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "var a = Tensor([10.0, 6, -4])\n",
    "var b = Tensor([2.0, 8, 7])\n",
    "(a,b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[12.0, 14.0, 3.0]\n"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a + b"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Comparison operators (>,<,==,!=,...) are `true` if all the elements of the tensors satisfy the comparison. Elementwise versions have the . prefix: .>, .<, .== ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "false\n"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a < b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[false, true, true]\n"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a .< b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]\n"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "var m = Tensor([1.0, 2, 3, 4, 5, 6, 7, 8, 9]).reshaped(to: [3,3])\n",
    "m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[2.0, 4.0, 6.0], [8.0, 10.0, 12.0], [14.0, 16.0, 18.0]]\n"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "2 * m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "16.881943016134134\n"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sqrt((m * m).sum().scalar!)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Elementwise matmul"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "func elementWiseMatmul(_ a:Tensor<Float>, _ b:Tensor<Float>) -> Tensor<Float>{\n",
    "    var res = Tensor<Float>(zeros: [a.shape[1], b.shape[0]])\n",
    "    for i in 0..<a.shape[0]{\n",
    "        for j in 0..<b.shape[1]{\n",
    "            res[i][j] = (a[i] * b.slice(lowerBounds: [0,j], upperBounds: [a.shape[1],j+1]).squeezingShape(at: 1)).sum()\n",
    "        }\n",
    "    }\n",
    "    return res\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "let res = elementWiseMatmul(m1, m2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "386.279412 ms\r\n"
     ]
    }
   ],
   "source": [
    "time() { let _ = elementWiseMatmul(m1, m2)}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Broadcasting"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Broadcasting with a scalar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "var a = Tensor([10.0, 6.0, -4.0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[true, true, false]\r\n"
     ]
    }
   ],
   "source": [
    "print(a .> 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In Tensorflow the operator `>` between tensors will return `true` if all the elements of the first tensor are greater than the ones of the second tensor. `.>` makes the memberwise comparison."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[11.0, 7.0, -3.0]\r\n"
     ]
    }
   ],
   "source": [
    "print(a+1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[2.0, 4.0, 6.0], [8.0, 10.0, 12.0], [14.0, 16.0, 18.0]]\n"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "2 * m"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Broadcasting a vector with a matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "let c = Tensor([10.0,20.0,30.0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By default, broadcasting is done by adding 1 dimensions to the beginning until dimensions of both objects match."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[11.0, 22.0, 33.0], [14.0, 25.0, 36.0], [17.0, 28.0, 39.0]]\n"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m + c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[11.0, 22.0, 33.0], [14.0, 25.0, 36.0], [17.0, 28.0, 39.0]]\n"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "c + m"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To broadcast on the other dimensions, one has to use `expandingShape` to add the dimension."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[11.0, 12.0, 13.0], [24.0, 25.0, 26.0], [37.0, 38.0, 39.0]]\n"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m + c.expandingShape(at: 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[10.0], [20.0], [30.0]]\n"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "c.expandingShape(at: 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Matmult with broadcasting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "func broadcastMatmult(_ a:Tensor<Float>, _ b:Tensor<Float>) -> Tensor<Float>{\n",
    "    var res = Tensor<Float>(zeros: [a.shape[0], b.shape[1]])\n",
    "    for i in 0..<a.shape[0]{\n",
    "        res[i] = (a[i].expandingShape(at: 1) * b).sum(squeezingAxes: 0)\n",
    "    }\n",
    "    return res\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "let res = broadcastMatmult(m1, m2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.9344620300000006 ms\r\n"
     ]
    }
   ],
   "source": [
    "time(repeating: 100) { let _ = broadcastMatmult(m1, m2)}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A faster way is to concatenate the columns instead of filling the matrix result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "func fastBcMatmul(_ a:Tensor<Float>, _ b:Tensor<Float>) -> Tensor<Float>{\n",
    "    var res: [Tensor<Float>] = []\n",
    "    for i in 0..<a.shape[0]{\n",
    "        res.append((a[i].expandingShape(at: 1) * b).sum(alongAxes: 0))\n",
    "    }\n",
    "    return Tensor<Float>(concatenating: res)\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.7782323200000002 ms\r\n"
     ]
    }
   ],
   "source": [
    "time(repeating: 100) { let _ = fastBcMatmul(m1, m2)}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Broadcasting rules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "▿ TensorShape\n",
       "  ▿ dimensions : 2 elements\n",
       "    - 0 : 1\n",
       "    - 1 : 3\n"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "c.expandingShape(at: 0).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "▿ TensorShape\n",
       "  ▿ dimensions : 2 elements\n",
       "    - 0 : 3\n",
       "    - 1 : 1\n"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "c.expandingShape(at: 1).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[100.0, 200.0, 300.0], [200.0, 400.0, 600.0], [300.0, 600.0, 900.0]]\n"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "c.expandingShape(at: 0) * c.expandingShape(at: 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[false, true, true], [false, false, true], [false, false, false]]\n"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "c.expandingShape(at: 0) .> c.expandingShape(at: 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Tensorflow op"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.024246269999999986 ms\r\n"
     ]
    }
   ],
   "source": [
    "time(repeating: 100) { let _ = matmul(m1, m2)}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Export"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "notebookToScript(fname: (Path.cwd / \"01_matmul.ipynb\").string)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Swift",
   "language": "swift",
   "name": "swift"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}