Chapter13 - Intro to Automatic Differentiation - Let's Build A Deep Learning Framework.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 1: Introduction to Tensors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tensor_v1([1, 2, 3, 4, 5])[2, 4, 6, 8, 10]"
     ]
    }
   ],
   "source": [
    "import Base:println,+\n",
    "\n",
    "abstract type Tensor end\n",
    "\n",
    "mutable struct Tensor_v1 <: Tensor\n",
    "    data \n",
    "end\n",
    "\n",
    "+(a::Tensor, b::Tensor) = a.data + b.data\n",
    "\n",
    "println(t::Tensor) = println(t.data)\n",
    "    \n",
    "x = Tensor_v1([1,2,3,4,5])\n",
    "print(x)\n",
    "\n",
    "y = x + x\n",
    "print(y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 2: Introduction to Autograd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "import Base:println, +, show\n",
    "\n",
    "mutable struct Tensor_v2 <: Tensor\n",
    "    data\n",
    "    creators\n",
    "    creation_op\n",
    "    grad \n",
    "    Tensor_v2(data; creators=nothing, creation_op = nothing) = \n",
    "    new(data, creators, creation_op)\n",
    "end\n",
    "\n",
    "function backward(t::Tensor, grad)\n",
    "    t.grad = grad\n",
    "    \n",
    "    if t.creation_op == \"add\"\n",
    "        backward(t.creators[1], grad)\n",
    "        backward(t.creators[2], grad)\n",
    "    end\n",
    "end\n",
    "\n",
    "+(a::Tensor, b::Tensor) = Tensor_v2(a.data + b.data; creators=[a,b], creation_op=\"add\")\n",
    "\n",
    "println(t::Tensor) = println(t.data)\n",
    "println(t::Array{Tensor_v2,1}) = println([i.data for i in t])\n",
    "show(io::IO,m::MIME\"text/plain\",a::Tensor) = show(io,m,a.data)\n",
    "    \n",
    "x = Tensor_v2([1,2,3,4,5])\n",
    "y = Tensor_v2([2,2,2,2,2])\n",
    "\n",
    "z = x + y\n",
    "backward(z, Tensor_v2([1,1,1,1,1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1, 1, 1, 1, 1]\n",
      "[1, 1, 1, 1, 1]\n",
      "[[1, 2, 3, 4, 5], [2, 2, 2, 2, 2]]\n",
      "add\n"
     ]
    }
   ],
   "source": [
    "println(x.grad)\n",
    "println(y.grad)\n",
    "println(z.creators)\n",
    "println(z.creation_op)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1, 1, 1, 1, 1]\n"
     ]
    }
   ],
   "source": [
    "a = Tensor_v2([1,2,3,4,5])\n",
    "b = Tensor_v2([2,2,2,2,2])\n",
    "c = Tensor_v2([5,4,3,2,1])\n",
    "d = Tensor_v2([-1,-2,-3,-4,-5])\n",
    "\n",
    "e = a + b\n",
    "f = c + d\n",
    "g = e + f\n",
    "\n",
    "backward(g, Tensor_v2([1,1,1,1,1]))\n",
    "\n",
    "println(a.grad)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 3: Tensors That Are Used Multiple Times"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "false"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a = Tensor_v2([1,2,3,4,5])\n",
    "b = Tensor_v2([2,2,2,2,2])\n",
    "c = Tensor_v2([5,4,3,2,1])\n",
    "\n",
    "d = a + b\n",
    "e = b + c\n",
    "f = d + e\n",
    "backward(f, Tensor_v2([1,1,1,1,1]))\n",
    "\n",
    "b.grad.data == [2,2,2,2,2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5-element Array{Int64,1}:\n",
       " 1\n",
       " 1\n",
       " 1\n",
       " 1\n",
       " 1"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b.grad.data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 4: Upgrading Autograd to Support Multiple Tensors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "true\n"
     ]
    }
   ],
   "source": [
    "using Random\n",
    "import Base:+,println\n",
    "\n",
    "mutable struct Tensor_v3 <: Tensor\n",
    "    data\n",
    "    autograd\n",
    "    creators\n",
    "    creation_op\n",
    "    id\n",
    "    children\n",
    "    grad \n",
    "    function Tensor_v3(data; autograd=false, creators=nothing, creation_op = nothing, id=nothing)\n",
    "        if isnothing(id)\n",
    "            id = rand(1:100000)\n",
    "        end\n",
    "        T = new(data, autograd, creators, creation_op, id)\n",
    "        T.children = Dict()\n",
    "        T.grad = nothing\n",
    "        \n",
    "        if !(isnothing(creators))\n",
    "            for c in creators\n",
    "                if haskey(c.children, T.id)\n",
    "                    c.children[T.id] += 1\n",
    "                else\n",
    "                    c.children[T.id] = 1\n",
    "                end\n",
    "            end\n",
    "        end\n",
    "        return T\n",
    "    end\n",
    "end\n",
    "\n",
    "function all_children_grads_accounted_for(t::Tensor)\n",
    "    for (id, cnt) in t.children\n",
    "        if (cnt != 0)\n",
    "            return false\n",
    "        end\n",
    "    end\n",
    "    return true\n",
    "end\n",
    "\n",
    "function backward(t::Tensor_v3, grad=nothing, grad_origin=nothing)\n",
    "    if t.autograd\n",
    "        grad = Tensor_v3(ones(size(t.data)))\n",
    "    \n",
    "        if !(isnothing(grad_origin))\n",
    "            if t.children[grad_origin.id] == 0\n",
    "                throw(\"cannot backprop more than once\")\n",
    "            else\n",
    "                t.children[grad_origin.id] -= 1\n",
    "            end\n",
    "        end\n",
    "        \n",
    "        if isnothing(t.grad)\n",
    "            t.grad = grad\n",
    "        else\n",
    "            t.grad += grad\n",
    "        end\n",
    "        \n",
    "        # grads must not have grads of their own\n",
    "        @assert !grad.autograd\n",
    "        \n",
    "        # only continue backpropping if there's something to\n",
    "        # backprop into and if all gradients (from children)\n",
    "        # are accounted for override waiting for children if\n",
    "        # \"backprop\" was called on this variable directly\n",
    "        \n",
    "        if (!isnothing(t.creators) && (all_children_grads_accounted_for(t) || isnothing(grad_origin)))\n",
    "            if t.creation_op == \"add\"\n",
    "                backward(t.creators[1], t.grad, t)\n",
    "                backward(t.creators[2], t.grad, t)\n",
    "            end\n",
    "        end\n",
    "    end\n",
    "end\n",
    "\n",
    "function +(a::Tensor_v3, b::Tensor_v3)\n",
    "    if (a.autograd && b.autograd)\n",
    "        return Tensor_v3(a.data .+ b.data; autograd=true, creators=[a,b], creation_op = \"add\")\n",
    "    end\n",
    "    return Tensor_v3(a.data+b.data)\n",
    "end\n",
    "\n",
    "println(t::Tensor_v3) = println(t.data)\n",
    "\n",
    "a = Tensor_v3([1,2,3,4,5]; autograd=true)\n",
    "b = Tensor_v3([2,2,2,2,2]; autograd=true)\n",
    "c = Tensor_v3([5,4,3,2,1]; autograd=true)\n",
    "\n",
    "d = a + b\n",
    "e = b + c\n",
    "f = d + e\n",
    "\n",
    "backward(f, Tensor_v3([1,1,1,1,1]))\n",
    "\n",
    "println(b.grad.data == [2,2,2,2,2])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 5: Add Support for Negation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "true"
     ]
    }
   ],
   "source": [
    "using Random\n",
    "import Base:+,-,println\n",
    "\n",
    "\n",
    "function backward(t::Tensor_v3, grad=nothing, grad_origin=nothing)\n",
    "    if t.autograd\n",
    "        if isnothing(grad)\n",
    "            grad = Tensor_v3(ones(size(t.data)))\n",
    "        end\n",
    "    \n",
    "        if !(isnothing(grad_origin))\n",
    "            if t.children[grad_origin.id] == 0\n",
    "                throw(\"cannot backprop more than once\")\n",
    "            else\n",
    "                t.children[grad_origin.id] -= 1\n",
    "            end\n",
    "        end\n",
    "        \n",
    "        if isnothing(t.grad)\n",
    "            t.grad = grad\n",
    "        else\n",
    "            t.grad += grad\n",
    "        end\n",
    "        \n",
    "        # grads must not have grads of their own\n",
    "        @assert !grad.autograd\n",
    "        \n",
    "        # only continue backpropping if there's something to\n",
    "        # backprop into and if all gradients (from children)\n",
    "        # are accounted for override waiting for children if\n",
    "        # \"backprop\" was called on this variable directly\n",
    "        \n",
    "        if (!isnothing(t.creators) && (all_children_grads_accounted_for(t) || isnothing(grad_origin)))\n",
    "            if t.creation_op == \"add\"\n",
    "                backward(t.creators[1], t.grad, t)\n",
    "                backward(t.creators[2], t.grad, t)\n",
    "            end\n",
    "            if t.creation_op == \"neg\"\n",
    "                backward(t.creators[1], -t.grad)\n",
    "            end\n",
    "        end\n",
    "    end\n",
    "end\n",
    "\n",
    "function -(a::Tensor_v3)\n",
    "    if (a.autograd)\n",
    "        return Tensor_v3(a.data .* -1; autograd=true, creators=[a], creation_op = \"neg\")\n",
    "    end\n",
    "    return Tensor_v3(a.data .* -1)\n",
    "end\n",
    "\n",
    "a = Tensor_v3([1,2,3,4,5]; autograd=true)\n",
    "b = Tensor_v3([2,2,2,2,2]; autograd=true)\n",
    "c = Tensor_v3([5,4,3,2,1]; autograd=true)\n",
    "\n",
    "d = a + (-b)\n",
    "e = (-b) + c\n",
    "f = d + e\n",
    "\n",
    "backward(f, Tensor_v3([1,1,1,1,1]))\n",
    "\n",
    "print(b.grad.data == [-2,-2,-2,-2,-2])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 6: Add Support for Additional Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Bool[1, 1, 1, 1, 1]"
     ]
    }
   ],
   "source": [
    "using Random\n",
    "import Base:+,-,*,println, sum, broadcasted, size, adjoint, show, dropdims\n",
    "\n",
    "\n",
    "function backward(t::Tensor_v3, grad=nothing, grad_origin=nothing)\n",
    "    if t.autograd\n",
    "        if isnothing(grad)\n",
    "            grad = Tensor_v3(ones(size(t.data)))\n",
    "        end\n",
    "    \n",
    "        if !(isnothing(grad_origin))\n",
    "            if t.children[grad_origin.id] == 0\n",
    "                throw(\"cannot backprop more than once\")\n",
    "            else\n",
    "                t.children[grad_origin.id] -= 1\n",
    "            end\n",
    "        end\n",
    "        \n",
    "        if isnothing(t.grad)\n",
    "            t.grad = grad\n",
    "        else\n",
    "            t.grad += grad\n",
    "        end\n",
    "        \n",
    "        # grads must not have grads of their own\n",
    "        @assert !grad.autograd\n",
    "        \n",
    "        # only continue backpropping if there's something to\n",
    "        # backprop into and if all gradients (from children)\n",
    "        # are accounted for override waiting for children if\n",
    "        # \"backprop\" was called on this variable directly\n",
    "        \n",
    "        if (!isnothing(t.creators) && (all_children_grads_accounted_for(t) || isnothing(grad_origin)))\n",
    "            if t.creation_op == \"add\"\n",
    "                backward(t.creators[1], t.grad, t)\n",
    "                backward(t.creators[2], t.grad, t)\n",
    "            end\n",
    "            \n",
    "            if t.creation_op == \"sub\"\n",
    "                backward(t.creators[1], t.grad, t)\n",
    "                backward(t.creators[2], -t.grad, t)\n",
    "            end\n",
    "            \n",
    "            if t.creation_op == \"mul\"\n",
    "                new_ = t.grad .* t.creators[2]\n",
    "                backward(t.creators[1], new_, t)\n",
    "                new_ = t.grad .* t.creators[1]\n",
    "                backward(t.creators[2], new_, t)\n",
    "            end\n",
    "            \n",
    "            if t.creation_op == \"mm\"\n",
    "                c1 = t.creators[1]\n",
    "                c2 = t.creators[2]\n",
    "                new_ =  t.grad * c2' ################\n",
    "                backward(c1, new_)\n",
    "                new_ = c1' * t.grad\n",
    "                backward(c2, new_)\n",
    "            end\n",
    "                  \n",
    "            if t.creation_op == \"transpose\"\n",
    "                backward(t.creators[1], t.grad')\n",
    "            end\n",
    "            \n",
    "            if occursin(\"sum\", t.creation_op)\n",
    "                dim = parse(Int, split(t.creation_op, \"_\")[2])\n",
    "                backward(t.creators[1], expand(t.grad, dim, size(t.creators[1].data)[dim]))\n",
    "            end\n",
    "            \n",
    "            if occursin(\"expand\", t.creation_op)\n",
    "                dim = parse(Int, split(t.creation_op, \"_\")[2])\n",
    "                ndims_cr = ndims(t.creators[1].data)\n",
    "                backward(t.creators[1], dropdims(sum(t.grad;dims=dim);dims=dim, ndims_cr=ndims_cr))\n",
    "            end              \n",
    "            \n",
    "            if t.creation_op == \"neg\"\n",
    "                backward(t.creators[1], -t.grad)\n",
    "            end\n",
    "        end\n",
    "    end\n",
    "end\n",
    "\n",
    "size(a::Tensor_v3) = size(a.data)\n",
    "\n",
    "function -(a::Tensor_v3, b::Tensor_v3)\n",
    "    if (a.autograd && b.autograd)\n",
    "        return Tensor_v3(a.data - b.data; autograd=true, creators=[a,b], creation_op = \"sub\")\n",
    "    end\n",
    "    return Tensor_v3(a.data-b.data)\n",
    "end\n",
    "\n",
    "#element-wise multiplication\n",
    "function broadcasted(f::typeof(*), a::Tensor_v3, b::Tensor_v3)\n",
    "    new_data = zeros(size(a.data))\n",
    "    for i=1:length(new_data)\n",
    "        new_data[i] = f(a.data[i] ,b.data[i])\n",
    "    end\n",
    "    if (a.autograd && b.autograd)\n",
    "        return Tensor_v3(new_data; autograd=true, creators=[a,b], creation_op =\"mul\")\n",
    "    end\n",
    "    return Tensor_v3(new_data)\n",
    "end\n",
    "\n",
    "function broadcasted(f::typeof(-), a::Tensor_v3, b::Tensor_v3)\n",
    "    new_data = zeros(size(a.data))\n",
    "    for i=1:length(new_data)\n",
    "        new_data[i] = -(a.data[i] ,b.data[i])\n",
    "    end\n",
    "    if (a.autograd && b.autograd)\n",
    "        return Tensor_v3(new_data; autograd=true, creators=[a,b], creation_op =\"sub\")\n",
    "    end\n",
    "    return Tensor_v3(new_data)\n",
    "end\n",
    "\n",
    "function sum(a::Tensor_v3; dims=dims)\n",
    "    new_ = dropdims(sum(a.data ;dims=dims), dims = tuple(findall(size(a) .== 1)...))\n",
    "    if (a.autograd)\n",
    "        return Tensor_v3(new_; autograd=true, creators=[a], creation_op = \"sum_\"*string(dims))\n",
    "    end\n",
    "    return Tensor_v3(new_)\n",
    "end\n",
    "\n",
    "function dropdims(a::Tensor_v3;dims=dims,ndims_cr=ndims_cr)\n",
    "    if ndims(a.data) == ndims_cr\n",
    "        return a\n",
    "    end\n",
    "    if (a.autograd)\n",
    "        return Tensor_v3(dropdims(a.data ;dims=dims); autograd=true, creators=[a], creation_op = \"dropdims\")\n",
    "    end\n",
    "    return Tensor_v3(dropdims(a.data ;dims=dims))\n",
    "end\n",
    "\n",
    "function expand(a::Tensor_v3, dim, copies)\n",
    "    sz = size(a)\n",
    "    rep = ntuple(d->d==dim ? copies : 1, length(sz)+1)\n",
    "    new_size = ntuple(d->d<dim ? sz[d] : d == dim ? 1 : sz[d-1], length(sz)+1)\n",
    "    new_data =  repeat(reshape(a.data, new_size), outer=rep)\n",
    "    if (a.autograd)\n",
    "        return Tensor_v3(new_data; autograd=true, creators=[a], creation_op = \"expand_\"*string(dim))\n",
    "    end\n",
    "    return Tensor_v3(new_data)\n",
    "end\n",
    "\n",
    "#transpose\n",
    "function adjoint(a::Tensor_v3)\n",
    "    if (a.autograd)\n",
    "        return Tensor_v3(a.data';autograd=true, creators=[a], creation_op = \"transpose\")\n",
    "    end\n",
    "    return Tensor_v3(a.data')\n",
    "end\n",
    "\n",
    "#matrix multiply \n",
    "function *(a::Tensor_v3, b::Tensor_v3)\n",
    "    if (a.autograd && b.autograd)\n",
    "        return Tensor_v3(a.data * b.data; autograd=true, creators=[a,b], creation_op = \"mm\")\n",
    "    end\n",
    "    return Tensor_v3(a.data * b.data)\n",
    "end\n",
    "\n",
    "a = Tensor_v3([1,2,3,4,5]; autograd=true)\n",
    "b = Tensor_v3([2,2,2,2,2]; autograd=true)\n",
    "c = Tensor_v3([5,4,3,2,1]; autograd=true)\n",
    "\n",
    "d = a + (-b)\n",
    "e = (-b) + c\n",
    "f = d + e\n",
    "\n",
    "backward(f, Tensor_v3([1,1,1,1,1]))\n",
    "\n",
    "print(b.grad.data .== [-2,-2,-2,-2,-2])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# A few Notes on Sum and Expand"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2×3 Array{Int64,2}:\n",
       " 1  2  3\n",
       " 4  5  6"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = Tensor_v3([1 2 3;4 5 6];autograd=true)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2×1 Array{Int64,2}:\n",
       "  6\n",
       " 15"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sum(x;dims=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1×3 Array{Int64,2}:\n",
       " 5  7  9"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sum(x;dims=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2×3×4 Array{Int64,3}:\n",
       "[:, :, 1] =\n",
       " 1  2  3\n",
       " 4  5  6\n",
       "\n",
       "[:, :, 2] =\n",
       " 1  2  3\n",
       " 4  5  6\n",
       "\n",
       "[:, :, 3] =\n",
       " 1  2  3\n",
       " 4  5  6\n",
       "\n",
       "[:, :, 4] =\n",
       " 1  2  3\n",
       " 4  5  6"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "expand(x,3,4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 7: Use Autograd to Train a Neural Network"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Previously we would train a model like this"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.319677517363771\n",
      "0.645370221071993\n",
      "0.48046558059136113\n",
      "0.42673209144608215\n",
      "0.3930618130835104\n",
      "0.363594896295545\n",
      "0.33605797739677096\n",
      "0.3100545140899324\n",
      "0.28544274299374384\n",
      "0.2621351937537131\n"
     ]
    }
   ],
   "source": [
    "using Random: seed!\n",
    "seed!(0)\n",
    "\n",
    "data = [ 0 0; 0 1; 1 0; 1 1;]\n",
    "target = [0; 1; 0; 1]\n",
    "\n",
    "weights_0_1 = rand(2,3)\n",
    "weights_1_2 = rand(3,1)\n",
    "\n",
    "for i=1:10\n",
    "    \n",
    "#     # Predict\n",
    "    layer_1 = data * weights_0_1\n",
    "    layer_2 = layer_1 * weights_1_2\n",
    "    \n",
    "#     # Compare\n",
    "    diff = (layer_2 - target)\n",
    "    sqdiff = (diff .* diff)\n",
    "    loss = sum(sqdiff;dims=1) # mean squared error loss\n",
    "\n",
    "#     # Learn: this is the backpropagation piece\n",
    "    layer_1_grad = diff * weights_1_2'\n",
    "    weight_1_2_update = layer_1' * diff\n",
    "    weight_0_1_update = data' * layer_1_grad\n",
    "    \n",
    "    weights_1_2 .-= weight_1_2_update .* 0.1\n",
    "    weights_0_1 .-= weight_0_1_update .* 0.1\n",
    "    println(loss[1])\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.6764031397217589]\n",
      "[0.23583373606478314]\n",
      "[0.09934872031663397]\n",
      "[0.039683547503183744]\n",
      "[0.015016987627295843]\n",
      "[0.0054452677736085454]\n",
      "[0.0019163531702008321]\n",
      "[0.0006614488249024953]\n",
      "[0.0002255690501204964]\n",
      "[7.636649139276776e-5]\n"
     ]
    }
   ],
   "source": [
    "using Random: seed!\n",
    "seed!(0)\n",
    "\n",
    "data = Tensor_v3([ 0 0; 0 1; 1 0; 1 1;], autograd=true)\n",
    "target = Tensor_v3([0; 1; 0; 1], autograd=true)\n",
    "\n",
    "w = []\n",
    "push!(w, Tensor_v3(rand(2,3), autograd=true))\n",
    "push!(w, Tensor_v3(rand(3,1), autograd=true))\n",
    "\n",
    "for i=1:10\n",
    "\n",
    "#     # Predict\n",
    "    pred_1 = data * w[1]\n",
    "    pred_2 = pred_1 * w[2]\n",
    "    diff_1 = pred_2 .- target\n",
    "    diff_2 = diff_1 .* diff_1\n",
    "    \n",
    "#     # Compare\n",
    "    loss = sum(diff_2;dims=1)\n",
    "    \n",
    "#     # Learn\n",
    "    backward(loss, Tensor_v3(ones(Float32, size(loss.data))))\n",
    "\n",
    "    for w_ in w\n",
    "        w_.data .-= w_.grad.data .* 0.1\n",
    "        w_.grad.data .*= 0\n",
    "    end\n",
    "\n",
    "    println(loss)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 8: Adding Automatic Optimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "step (generic function with 2 methods)"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mutable struct SGD\n",
    "    parameters\n",
    "    alpha\n",
    "    SGD(parameters, alpha) = new(parameters, alpha)\n",
    "end\n",
    "\n",
    "function zero!(opt::SGD)\n",
    "    for p in opt.parameters\n",
    "        p.grad.data .*= 0.0\n",
    "    end\n",
    "end\n",
    "\n",
    "function step(opt::SGD, zero=true)\n",
    "    for p in opt.parameters\n",
    "        p.data -= (p.grad.data .* opt.alpha)\n",
    "        if zero\n",
    "            p.grad.data .*= 0.0\n",
    "        end\n",
    "    end\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.6764031397217589]\n",
      "[0.23583373606478314]\n",
      "[0.09934872031663397]\n",
      "[0.039683547503183744]\n",
      "[0.015016987627295843]\n",
      "[0.0054452677736085454]\n",
      "[0.0019163531702008321]\n",
      "[0.0006614488249024953]\n",
      "[0.0002255690501204964]\n",
      "[7.636649139276776e-5]\n"
     ]
    }
   ],
   "source": [
    "using Random: seed!\n",
    "seed!(0)\n",
    "\n",
    "data = Tensor_v3([ 0 0; 0 1; 1 0; 1 1;], autograd=true)\n",
    "target = Tensor_v3([0; 1; 0; 1], autograd=true)\n",
    "\n",
    "w = []\n",
    "push!(w, Tensor_v3(rand(2,3), autograd=true))\n",
    "push!(w, Tensor_v3(rand(3,1), autograd=true))\n",
    "\n",
    "opt = SGD(w, 0.1)\n",
    "\n",
    "for i=1:10\n",
    "\n",
    "#     # Predict\n",
    "    pred_1 = data * w[1]\n",
    "    pred_2 = pred_1 * w[2]\n",
    "    diff_1 = pred_2 .- target\n",
    "    diff_2 = diff_1 .* diff_1\n",
    "    \n",
    "#     # Compare\n",
    "    loss = sum(diff_2;dims=1)\n",
    "    \n",
    "#     # Learn\n",
    "    backward(loss, Tensor_v3(ones(Float32, size(loss.data))))\n",
    "\n",
    "    step(opt)\n",
    "\n",
    "    println(loss)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 9: Adding Support for Layer Types"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "forward (generic function with 1 method)"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "abstract type Layer end\n",
    "\n",
    "function get_parameters(l::Layer)\n",
    "    return l.parameters\n",
    "end\n",
    "\n",
    "mutable struct Linear <: Layer\n",
    "    W\n",
    "    b\n",
    "    parameters\n",
    "    \n",
    "    function Linear(n_inputs, n_outputs)\n",
    "        W = Tensor_v3(randn(n_outputs, n_inputs) .* sqrt(1.0/n_inputs), autograd=true)\n",
    "        b = Tensor_v3(zeros(n_outputs), autograd=true)\n",
    "        parameters = [W,b]\n",
    "        return new(W,b,parameters)\n",
    "    end\n",
    "end\n",
    "\n",
    "function forward(l::Linear, input)\n",
    "    return (l.W * input)  + expand(l.b,2,size(input.data, 2))\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 10: Layers Which Contain Layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2.097122200876772]\n",
      "[1.090072584706214]\n",
      "[0.8848315119040038]\n",
      "[0.7588875978275563]\n",
      "[0.6498970345986242]\n",
      "[0.5416092338969248]\n",
      "[0.43359472039645397]\n",
      "[0.32996001398541697]\n",
      "[0.23650275816675004]\n",
      "[0.1586239346924714]\n"
     ]
    }
   ],
   "source": [
    "mutable struct Sequential <: Layer\n",
    "    layers\n",
    "    function Sequential(layers)\n",
    "        return new(layers)\n",
    "    end\n",
    "end\n",
    "\n",
    "function add(s::Sequential, layer)\n",
    "    push!(s.layers, layer)\n",
    "end\n",
    "\n",
    "function forward(s::Sequential, input)\n",
    "    for layer in s.layers\n",
    "        input = forward(layer, input)\n",
    "    end\n",
    "    return input\n",
    "end\n",
    "\n",
    "function get_parameters(s::Sequential)\n",
    "    parameters = [get_parameters(layer) for layer in s.layers]\n",
    "    return collect(Iterators.flatten(parameters))\n",
    "end\n",
    "\n",
    "using Random: seed!; seed!(0)\n",
    "data = Tensor_v3([ 0  0  1  1;0  1  0  1], autograd=true)\n",
    "target = Tensor_v3([0 1 0 1], autograd=true)\n",
    "\n",
    "model = Sequential([Linear(2,3), Linear(3,1)])\n",
    "\n",
    "optim = SGD(get_parameters(model),0.1)\n",
    "\n",
    "for i=1:10\n",
    "    \n",
    "    # Predict\n",
    "    pred = forward(model, data)\n",
    "    \n",
    "    # Compare\n",
    "    loss = sum((pred - target) .* (pred - target);dims=2)\n",
    "    \n",
    "    # Learn\n",
    "    backward(loss, Tensor_v3(ones(Float32, size(loss.data))))\n",
    "    step(optim)\n",
    "    println(loss)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 11: Loss Function Layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2.097122200876772]\n",
      "[1.090072584706214]\n",
      "[0.8848315119040038]\n",
      "[0.7588875978275563]\n",
      "[0.6498970345986242]\n",
      "[0.5416092338969248]\n",
      "[0.43359472039645397]\n",
      "[0.32996001398541697]\n",
      "[0.23650275816675004]\n",
      "[0.1586239346924714]\n"
     ]
    }
   ],
   "source": [
    "struct MSELoss <: Layer\n",
    "    MSELoss() = new()\n",
    "end\n",
    "\n",
    "function forward(l::MSELoss, pred, target)\n",
    "    return sum((pred - target) .* (pred - target);dims=2)\n",
    "end\n",
    "\n",
    "using Random: seed!; seed!(0)\n",
    "data = Tensor_v3([ 0  0  1  1;0  1  0  1], autograd=true)\n",
    "target = Tensor_v3([0 1 0 1], autograd=true)\n",
    "\n",
    "model = Sequential([Linear(2,3), Linear(3,1)])\n",
    "criterion = MSELoss()\n",
    "\n",
    "optim = SGD(get_parameters(model),0.1)\n",
    "\n",
    "for i=1:10\n",
    "    \n",
    "    # Predict\n",
    "    pred = forward(model, data)\n",
    "    \n",
    "    # Compare\n",
    "    loss = forward(criterion,pred, target)\n",
    "    \n",
    "    # Learn\n",
    "    backward(loss, Tensor_v3(ones(Float32, size(loss.data))))\n",
    "    step(optim)\n",
    "    println(loss)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 12: Non-linearity Layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "forward (generic function with 7 methods)"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "using Random\n",
    "import Base:+,-,*,println, sum, broadcasted, size, adjoint, show, dropdims, tanh\n",
    "\n",
    "\n",
    "function backward(t::Tensor_v3, grad=nothing, grad_origin=nothing)\n",
    "    if t.autograd\n",
    "        if isnothing(grad)\n",
    "            grad = Tensor_v3(ones(size(t.data)))\n",
    "        end\n",
    "    \n",
    "        if !(isnothing(grad_origin))\n",
    "            if t.children[grad_origin.id] == 0\n",
    "                throw(\"cannot backprop more than once\")\n",
    "            else\n",
    "                t.children[grad_origin.id] -= 1\n",
    "            end\n",
    "        end\n",
    "        \n",
    "        if isnothing(t.grad)\n",
    "            t.grad = grad\n",
    "        else\n",
    "            t.grad += grad\n",
    "        end\n",
    "        \n",
    "        # grads must not have grads of their own\n",
    "        @assert !grad.autograd\n",
    "        \n",
    "        # only continue backpropping if there's something to\n",
    "        # backprop into and if all gradients (from children)\n",
    "        # are accounted for override waiting for children if\n",
    "        # \"backprop\" was called on this variable directly\n",
    "        \n",
    "        if (!isnothing(t.creators) && (all_children_grads_accounted_for(t) || isnothing(grad_origin)))\n",
    "            if t.creation_op == \"add\"\n",
    "                backward(t.creators[1], t.grad, t)\n",
    "                backward(t.creators[2], t.grad, t)\n",
    "            end\n",
    "            \n",
    "            if t.creation_op == \"sub\"\n",
    "                backward(t.creators[1], t.grad, t)\n",
    "                backward(t.creators[2], -t.grad, t)\n",
    "            end\n",
    "            \n",
    "            if t.creation_op == \"mul\"\n",
    "                new_ = t.grad .* t.creators[2]\n",
    "                backward(t.creators[1], new_, t)\n",
    "                new_ = t.grad .* t.creators[1]\n",
    "                backward(t.creators[2], new_, t)\n",
    "            end\n",
    "            \n",
    "            if t.creation_op == \"mm\"\n",
    "                c1 = t.creators[1]\n",
    "                c2 = t.creators[2]\n",
    "                new_ =  t.grad * c2' ################\n",
    "                backward(c1, new_)\n",
    "                new_ = c1' * t.grad\n",
    "                backward(c2, new_)\n",
    "            end\n",
    "                  \n",
    "            if t.creation_op == \"transpose\"\n",
    "                backward(t.creators[1], t.grad')\n",
    "            end\n",
    "            \n",
    "            if occursin(\"sum\", t.creation_op)\n",
    "                dim = parse(Int, split(t.creation_op, \"_\")[2])\n",
    "                backward(t.creators[1], expand(t.grad, dim, size(t.creators[1].data)[dim]))\n",
    "            end\n",
    "            \n",
    "            if occursin(\"expand\", t.creation_op)\n",
    "                dim = parse(Int, split(t.creation_op, \"_\")[2])\n",
    "                ndims_cr = ndims(t.creators[1].data)\n",
    "                backward(t.creators[1], dropdims(sum(t.grad;dims=dim);dims=dim, ndims_cr=ndims_cr))\n",
    "            end\n",
    "            \n",
    "            if t.creation_op == \"neg\"\n",
    "                backward(t.creators[1], -t.grad)\n",
    "            end\n",
    "            \n",
    "            if t.creation_op == \"sigmoid\"\n",
    "                ones_ = Tensor_v3(ones(size(t.grad.data)))\n",
    "                backward(t.creators[1], t.grad .* t .* (ones_ - t) )\n",
    "            end\n",
    "            \n",
    "            if t.creation_op == \"tanh\"\n",
    "                ones_ = Tensor_v3(ones(size(t.grad.data)))\n",
    "                backward(t.creators[1], t.grad .* (ones_ - (t .* t)))\n",
    "            end\n",
    "        end\n",
    "    end\n",
    "end\n",
    "\n",
    "σ(x) = 1/(1+exp(-x))        \n",
    "\n",
    "struct Tanh <: Layer\n",
    "    Tanh() = new()\n",
    "end\n",
    "\n",
    "struct Sigmoid <: Layer\n",
    "    Sigmoid() = new()\n",
    "end\n",
    "\n",
    "function get_parameters(l::Tanh)\n",
    "    return []\n",
    "end\n",
    "\n",
    "function get_parameters(l::Sigmoid)\n",
    "    return []\n",
    "end\n",
    "\n",
    "function forward(l::Sigmoid, a::Tensor_v3)\n",
    "    if a.autograd\n",
    "        return Tensor_v3(σ.(a.data); autograd=true, creators=[a], creation_op = \"sigmoid\")\n",
    "    end\n",
    "    return Tensor_v3(σ.(a.data))\n",
    "end\n",
    "        \n",
    "function forward(l::Tanh, a::Tensor_v3)\n",
    "    if a.autograd\n",
    "        return Tensor_v3(tanh.(a.data); autograd=true, creators=[a], creation_op = \"tanh\")\n",
    "    end\n",
    "    return Tensor_v3(tanh.(a.data))\n",
    "end   \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1.0640387281517036]\n",
      "[0.9685119577123646]\n",
      "[0.8989054438898241]\n",
      "[0.8006851582324737]\n",
      "[0.65271510312216]\n",
      "[0.46696476003616044]\n",
      "[0.2989261336464314]\n",
      "[0.18861277315979982]\n",
      "[0.12615799262916702]\n",
      "[0.09048263176049438]\n"
     ]
    }
   ],
   "source": [
    "using Random: seed!; seed!(0)\n",
    "data = Tensor_v3([ 0  0  1  1;0  1  0  1], autograd=true)\n",
    "target = Tensor_v3([0 1 0 1], autograd=true)\n",
    "\n",
    "model = Sequential([Linear(2,3), Tanh(), Linear(3,1), Sigmoid()])\n",
    "criterion = MSELoss()\n",
    "\n",
    "optim = SGD(get_parameters(model),1.0)\n",
    "\n",
    "for i=1:10\n",
    "    \n",
    "    # Predict\n",
    "    pred = forward(model, data)\n",
    "    \n",
    "    # Compare\n",
    "    loss = forward(criterion,pred, target)\n",
    "    \n",
    "    # Learn\n",
    "    backward(loss, Tensor_v3(ones(Float32, size(loss.data))))\n",
    "    step(optim)\n",
    "    println(loss)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 13: The Embedding Layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "mutable struct Embedding_v1 <: Layer\n",
    "    vocab_size\n",
    "    dim\n",
    "    weight\n",
    "    \n",
    "    # this random initialiation style is just a convention from word2vec\n",
    "    Embedding_v1(vocab_size, dim) = new(vocab_size, dim, (randn(dim, vocab_size) .- 0.5) ./ dim) \n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "forward (generic function with 10 methods)"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "using Random\n",
    "import Base:+,-,*,println, sum, broadcasted, size, adjoint, show, dropdims, tanh\n",
    "using Base.Iterators:partition, flatten\n",
    "\n",
    "mutable struct Tensor_v4 <: Tensor\n",
    "    data\n",
    "    autograd\n",
    "    creators\n",
    "    creation_op\n",
    "    id\n",
    "    children\n",
    "    grad \n",
    "    index_select_indices\n",
    "    function Tensor_v4(data; autograd=false, creators=nothing, creation_op = nothing, id=nothing)\n",
    "        if isnothing(id)\n",
    "            id = rand(1:100000)\n",
    "        end\n",
    "        T = new(data, autograd, creators, creation_op, id)\n",
    "        T.children = Dict()\n",
    "        T.grad = nothing\n",
    "        T.index_select_indices = nothing\n",
    "        \n",
    "        if !(isnothing(creators))\n",
    "            for c in creators\n",
    "                if haskey(c.children, T.id)\n",
    "                    c.children[T.id] += 1\n",
    "                else\n",
    "                    c.children[T.id] = 1\n",
    "                end\n",
    "            end\n",
    "        end\n",
    "        return T\n",
    "    end\n",
    "end\n",
    "\n",
    "function all_children_grads_accounted_for(t::Tensor_v4)\n",
    "    for (id, cnt) in t.children\n",
    "        if (cnt != 0)\n",
    "            return false\n",
    "        end\n",
    "    end\n",
    "    return true\n",
    "end\n",
    "\n",
    "function backward(t::Tensor_v4, grad=nothing, grad_origin=nothing)\n",
    "    if t.autograd\n",
    "        if isnothing(grad)\n",
    "            grad = Tensor_v4(ones(size(t.data)))\n",
    "        end\n",
    "    \n",
    "        if !(isnothing(grad_origin))\n",
    "            if t.children[grad_origin.id] == 0\n",
    "                throw(\"cannot backprop more than once\")\n",
    "            else\n",
    "                t.children[grad_origin.id] -= 1\n",
    "            end\n",
    "        end\n",
    "        \n",
    "        if isnothing(t.grad)\n",
    "            t.grad = grad\n",
    "        else\n",
    "            t.grad += grad\n",
    "        end\n",
    "        \n",
    "        # grads must not have grads of their own\n",
    "        @assert !grad.autograd\n",
    "        \n",
    "        # only continue backpropping if there's something to\n",
    "        # backprop into and if all gradients (from children)\n",
    "        # are accounted for override waiting for children if\n",
    "        # \"backprop\" was called on this variable directly\n",
    "        \n",
    "        if (!isnothing(t.creators) && (all_children_grads_accounted_for(t) || isnothing(grad_origin)))\n",
    "            if t.creation_op == \"add\"\n",
    "                backward(t.creators[1], t.grad, t)\n",
    "                backward(t.creators[2], t.grad, t)\n",
    "            end\n",
    "            \n",
    "            if t.creation_op == \"sub\"\n",
    "                backward(t.creators[1], t.grad, t)\n",
    "                backward(t.creators[2], -t.grad, t)\n",
    "            end\n",
    "            \n",
    "            if t.creation_op == \"mul\"\n",
    "                new_ = t.grad .* t.creators[2]\n",
    "                backward(t.creators[1], new_, t)\n",
    "                new_ = t.grad .* t.creators[1]\n",
    "                backward(t.creators[2], new_, t)\n",
    "            end\n",
    "            \n",
    "            if t.creation_op == \"mm\"\n",
    "                c1 = t.creators[1]\n",
    "                c2 = t.creators[2]\n",
    "                new_ =  t.grad * c2' ################\n",
    "                backward(c1, new_)\n",
    "                new_ = c1' * t.grad\n",
    "                backward(c2, new_)\n",
    "            end\n",
    "                  \n",
    "            if t.creation_op == \"transpose\"\n",
    "                backward(t.creators[1], t.grad')\n",
    "            end\n",
    "            \n",
    "            if occursin(\"sum\", t.creation_op)\n",
    "                dim = parse(Int, split(t.creation_op, \"_\")[2])\n",
    "                backward(t.creators[1], expand(t.grad, dim, size(t.creators[1].data)[dim]))\n",
    "            end\n",
    "            \n",
    "            if occursin(\"expand\", t.creation_op)\n",
    "                dim = parse(Int, split(t.creation_op, \"_\")[2])\n",
    "                ndims_cr = ndims(t.creators[1].data)\n",
    "                backward(t.creators[1], dropdims(sum(t.grad;dims=dim);dims=dim, ndims_cr=ndims_cr))\n",
    "            end\n",
    "            \n",
    "            if t.creation_op == \"neg\"\n",
    "                backward(t.creators[1], -t.grad)\n",
    "            end\n",
    "            \n",
    "            if t.creation_op == \"sigmoid\"\n",
    "                ones_ = Tensor_v4(ones(size(t.grad.data)))\n",
    "                backward(t.creators[1], t.grad .* t .* (ones_ - t) )\n",
    "            end\n",
    "            \n",
    "            if t.creation_op == \"tanh\"\n",
    "                ones_ = Tensor_v4(ones(size(t.grad.data)))\n",
    "                backward(t.creators[1], t.grad .* (ones_ - (t .* t)))\n",
    "            end\n",
    "            \n",
    "            if t.creation_op == \"index_select\"\n",
    "                new_grad = zeros(size(t.creators[1]))\n",
    "                indices = t.index_select_indices.data\n",
    "                major_chunks = partition(1:size(t.grad,2),length(indices))\n",
    "                grad_chunks = [t.grad.data[:,inds][:,j]  for(i,inds) in enumerate(major_chunks) for j=1:size(inds)[1]]\n",
    "    \n",
    "                for (i,ind) in enumerate(flatten(indices))\n",
    "                    new_grad[:,ind] +=  grad_chunks[i]\n",
    "                end\n",
    "                backward(t.creators[1], Tensor_v4(new_grad))\n",
    "            end\n",
    "        end\n",
    "    end\n",
    "end\n",
    "\n",
    "size(a::Tensor_v4) = size(a.data)\n",
    "size(a::Tensor_v4, ind::Int) = size(a.data, ind)\n",
    "\n",
    "function +(a::Tensor_v4, b::Tensor_v4)\n",
    "    if (a.autograd && b.autograd)\n",
    "        return Tensor_v4(a.data + b.data; autograd=true, creators=[a,b], creation_op = \"add\")\n",
    "    end\n",
    "    return Tensor_v4(a.data+b.data)\n",
    "end\n",
    "\n",
    "function -(a::Tensor_v4)\n",
    "    if (a.autograd)\n",
    "        return Tensor_v4(a.data .* -1; autograd=true, creators=[a], creation_op = \"neg\")\n",
    "    end\n",
    "    return Tensor_v4(a.data .* -1)\n",
    "end\n",
    "\n",
    "function -(a::Tensor_v4, b::Tensor_v4)\n",
    "    if (a.autograd && b.autograd)\n",
    "        return Tensor_v4(a.data - b.data; autograd=true, creators=[a,b], creation_op = \"sub\")\n",
    "    end\n",
    "    return Tensor_v4(a.data-b.data)\n",
    "end\n",
    "\n",
    "#element-wise multiplication\n",
    "function broadcasted(f::typeof(*), a::Tensor_v4, b::Tensor_v4)\n",
    "    new_data = zeros(size(a.data))\n",
    "    for i=1:length(new_data)\n",
    "        new_data[i] = f(a.data[i] ,b.data[i])\n",
    "    end\n",
    "    if (a.autograd && b.autograd)\n",
    "        return Tensor_v4(new_data; autograd=true, creators=[a,b], creation_op =\"mul\")\n",
    "    end\n",
    "    return Tensor_v4(new_data)\n",
    "end\n",
    "\n",
    "function broadcasted(f::typeof(-), a::Tensor_v4, b::Tensor_v4)\n",
    "    new_data = zeros(size(a.data))\n",
    "    for i=1:length(new_data)\n",
    "        new_data[i] = -(a.data[i] ,b.data[i])\n",
    "    end\n",
    "    if (a.autograd && b.autograd)\n",
    "        return Tensor_v4(new_data; autograd=true, creators=[a,b], creation_op =\"sub\")\n",
    "    end\n",
    "    return Tensor_v4(new_data)\n",
    "end\n",
    "\n",
    "function sum(a::Tensor_v4; dims=dims)\n",
    "    new_ = dropdims(sum(a.data ;dims=dims), dims = tuple(findall(size(a) .== 1)...))\n",
    "    if (a.autograd)\n",
    "        return Tensor_v4(new_; autograd=true, creators=[a], creation_op = \"sum_\"*string(dims))\n",
    "    end\n",
    "    return Tensor_v4(new_)\n",
    "end\n",
    "\n",
    "function dropdims(a::Tensor_v4;dims=dims,ndims_cr=ndims_cr)\n",
    "    if ndims(a.data) == ndims_cr\n",
    "        return a\n",
    "    end\n",
    "    if (a.autograd)\n",
    "        return Tensor_v4(dropdims(a.data ;dims=dims); autograd=true, creators=[a], creation_op = \"dropdims\")\n",
    "    end\n",
    "    return Tensor_v4(dropdims(a.data ;dims=dims))\n",
    "end\n",
    "\n",
    "function expand(a::Tensor_v4, dim, copies)\n",
    "    sz = size(a)\n",
    "    rep = ntuple(d->d==dim ? copies : 1, length(sz)+1)\n",
    "    new_size = ntuple(d->d<dim ? sz[d] : d == dim ? 1 : sz[d-1], length(sz)+1)\n",
    "    new_data =  repeat(reshape(a.data, new_size), outer=rep)\n",
    "    if (a.autograd)\n",
    "        return Tensor_v4(new_data; autograd=true, creators=[a], creation_op = \"expand_\"*string(dim))\n",
    "    end\n",
    "    return Tensor_v4(new_data)\n",
    "end\n",
    "\n",
    "#transpose\n",
    "function adjoint(a::Tensor_v4)\n",
    "    if (a.autograd)\n",
    "        return Tensor_v4(a.data';autograd=true, creators=[a], creation_op = \"transpose\")\n",
    "    end\n",
    "    return Tensor_v4(a.data')\n",
    "end\n",
    "\n",
    "#matrix multiply \n",
    "function *(a::Tensor_v4, b::Tensor_v4)\n",
    "    if (a.autograd && b.autograd)\n",
    "        return Tensor_v4(a.data * b.data; autograd=true, creators=[a,b], creation_op = \"mm\")\n",
    "    end\n",
    "    return Tensor_v4(a.data * b.data)\n",
    "end\n",
    "\n",
    "\n",
    "function index_select_helper(a::Array, indices)\n",
    "    return reduce(hcat,map(ind -> a[:,ind], indices))\n",
    "end\n",
    "\n",
    "function index_select(a::Tensor_v4, indices::Tensor_v4)\n",
    "    new_ = index_select_helper(a.data, indices.data)\n",
    "    if (a.autograd)\n",
    "        T = Tensor_v4(new_, autograd=true, creators=[a], creation_op = \"index_select\")\n",
    "        T.index_select_indices = indices\n",
    "        return T\n",
    "    end\n",
    "    return Tensor_v4(new_)\n",
    "end\n",
    "\n",
    "\n",
    "abstract type Layer end\n",
    "\n",
    "function get_parameters(l::Layer)\n",
    "    return l.parameters\n",
    "end\n",
    "\n",
    "mutable struct Linear <: Layer\n",
    "    W\n",
    "    b\n",
    "    parameters\n",
    "    \n",
    "    function Linear(n_inputs, n_outputs)\n",
    "        W = Tensor_v4(randn(n_outputs, n_inputs) .* sqrt(1.0/n_inputs), autograd=true)\n",
    "        b = Tensor_v4(zeros(n_outputs), autograd=true)\n",
    "        parameters = [W,b]\n",
    "        return new(W,b,parameters)\n",
    "    end\n",
    "end\n",
    "\n",
    "function forward(l::Linear, input)\n",
    "    return (l.W * input)  + expand(l.b,2,size(input.data, 2))\n",
    "end\n",
    "\n",
    "σ(x) = 1/(1+exp(-x))            \n",
    "\n",
    "println(t::Tensor_v4) = println(t.data)\n",
    "show(io::IO,m::MIME\"text/plain\",a::Tensor_v4) = show(io,m,a.data)\n",
    "\n",
    "struct Tanh <: Layer\n",
    "    Tanh() = new()\n",
    "end\n",
    "\n",
    "struct Sigmoid <: Layer\n",
    "    Sigmoid() = new()\n",
    "end\n",
    "\n",
    "function get_parameters(l::Tanh)\n",
    "    return []\n",
    "end\n",
    "\n",
    "function get_parameters(l::Sigmoid)\n",
    "    return []\n",
    "end\n",
    "\n",
    "function forward(l::Sigmoid, a::Tensor_v4)\n",
    "    if a.autograd\n",
    "        return Tensor_v4(σ.(a.data); autograd=true, creators=[a], creation_op = \"sigmoid\")\n",
    "    end\n",
    "    return Tensor_v4(σ.(a.data))\n",
    "end\n",
    "        \n",
    "function forward(l::Tanh, a::Tensor_v4)\n",
    "    if a.autograd\n",
    "        return Tensor_v4(tanh.(a.data); autograd=true, creators=[a], creation_op = \"tanh\")\n",
    "    end\n",
    "    return Tensor_v4(tanh.(a.data))\n",
    "end    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5×5 Array{Float64,2}:\n",
       " 0.0  1.0  2.0  2.0  1.0\n",
       " 0.0  1.0  2.0  2.0  1.0\n",
       " 0.0  1.0  2.0  2.0  1.0\n",
       " 0.0  1.0  2.0  2.0  1.0\n",
       " 0.0  1.0  2.0  2.0  1.0"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "using LinearAlgebra:I\n",
    "x = Tensor_v4(1.0* Matrix(I, 5, 5), autograd=true)\n",
    "backward(index_select(x, Tensor_v4([[2,3,4],[3,4,5]])))\n",
    "x.grad"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 15: The Embedding Layer (revisited)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "forward (generic function with 10 methods)"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mutable struct Embedding_v2 <: Layer\n",
    "    vocab_size\n",
    "    dim\n",
    "    weight\n",
    "    parameters\n",
    "    # this random initialiation style is just a convention from word2vec\n",
    "    function Embedding_v2(dim, vocab_size) \n",
    "        E = new(vocab_size, dim, Tensor_v4((randn(dim, vocab_size) .- 0.5) ./ dim; autograd=true))\n",
    "        E.parameters = [E.weight]\n",
    "        return E\n",
    "    end\n",
    "end\n",
    "\n",
    "function forward(E::Embedding_v2, indices)\n",
    "    return index_select(E.weight, indices)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.8728139692976189]\n",
      "[0.35474837968994166]\n",
      "[0.19046978251082405]\n",
      "[0.1201225262673748]\n",
      "[0.0849317092081909]\n",
      "[0.06457754793742479]\n",
      "[0.05154228148446023]\n",
      "[0.042573745189638766]\n",
      "[0.03607112398001075]\n",
      "[0.0311650999160659]\n"
     ]
    }
   ],
   "source": [
    "using Random:seed!;seed!(0)\n",
    "data = Tensor_v4([[1,2,1,2]], autograd=true)\n",
    "target = Tensor_v4([0 1 0 1], autograd=true)\n",
    "\n",
    "embed = Embedding_v2(3,5)\n",
    "model = Sequential([embed, Tanh(), Linear(3,1), Sigmoid()])\n",
    "criterion = MSELoss()\n",
    "\n",
    "\n",
    "optim = SGD(get_parameters(model),1.0)\n",
    "\n",
    "for i=1:10\n",
    "    \n",
    "    # Predict\n",
    "    pred = forward(model, data)\n",
    "    \n",
    "    # Compare\n",
    "    loss = forward(criterion,pred, target)\n",
    "    \n",
    "    # Learn\n",
    "    backward(loss, Tensor_v4(ones(Float32, size(loss.data))))\n",
    "    step(optim)\n",
    "    println(loss)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 16: The Cross Entropy Layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "forward (generic function with 13 methods)"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "using Random\n",
    "import Base:+,-,*,println, sum, broadcasted, size, adjoint, show, dropdims, tanh\n",
    "using Base.Iterators:partition, flatten\n",
    "\n",
    "mutable struct Tensor_v5\n",
    "    data\n",
    "    autograd\n",
    "    creators\n",
    "    creation_op\n",
    "    id\n",
    "    children\n",
    "    grad \n",
    "    index_select_indices\n",
    "    softmax_output\n",
    "    target_dist\n",
    "    \n",
    "    function Tensor_v5(data; autograd=false, creators=nothing, creation_op = nothing, id=nothing)\n",
    "        if isnothing(id)\n",
    "            id = rand(1:100000)\n",
    "        end\n",
    "        T = new(data, autograd, creators, creation_op, id)\n",
    "        T.children = Dict()\n",
    "        T.grad = nothing\n",
    "        T.index_select_indices = nothing\n",
    "        \n",
    "        if !(isnothing(creators))\n",
    "            for c in creators\n",
    "                if haskey(c.children, T.id)\n",
    "                    c.children[T.id] += 1\n",
    "                else\n",
    "                    c.children[T.id] = 1\n",
    "                end\n",
    "            end\n",
    "        end\n",
    "        return T\n",
    "    end\n",
    "end\n",
    "\n",
    "function all_children_grads_accounted_for(t::Tensor_v5)\n",
    "    for (id, cnt) in t.children\n",
    "        if (cnt != 0)\n",
    "            return false\n",
    "        end\n",
    "    end\n",
    "    return true\n",
    "end\n",
    "\n",
    "function backward(t::Tensor_v5, grad=nothing, grad_origin=nothing)\n",
    "    if t.autograd\n",
    "        if isnothing(grad)\n",
    "            grad = Tensor_v5(ones(size(t.data)))\n",
    "        end\n",
    "    \n",
    "        if !(isnothing(grad_origin))\n",
    "            if t.children[grad_origin.id] == 0\n",
    "                throw(\"cannot backprop more than once\")\n",
    "            else\n",
    "                t.children[grad_origin.id] -= 1\n",
    "            end\n",
    "        end\n",
    "        \n",
    "        if isnothing(t.grad)\n",
    "            t.grad = grad\n",
    "        else\n",
    "            t.grad += grad\n",
    "        end\n",
    "        \n",
    "        # grads must not have grads of their own\n",
    "        @assert !grad.autograd\n",
    "        \n",
    "        # only continue backpropping if there's something to\n",
    "        # backprop into and if all gradients (from children)\n",
    "        # are accounted for override waiting for children if\n",
    "        # \"backprop\" was called on this variable directly\n",
    "        \n",
    "        if (!isnothing(t.creators) && (all_children_grads_accounted_for(t) || isnothing(grad_origin)))\n",
    "            if t.creation_op == \"add\"\n",
    "                backward(t.creators[1], t.grad, t)\n",
    "                backward(t.creators[2], t.grad, t)\n",
    "            end\n",
    "            \n",
    "            if t.creation_op == \"sub\"\n",
    "                backward(t.creators[1], t.grad, t)\n",
    "                backward(t.creators[2], -t.grad, t)\n",
    "            end\n",
    "            \n",
    "            if t.creation_op == \"mul\"\n",
    "                new_ = t.grad .* t.creators[2]\n",
    "                backward(t.creators[1], new_, t)\n",
    "                new_ = t.grad .* t.creators[1]\n",
    "                backward(t.creators[2], new_, t)\n",
    "            end\n",
    "            \n",
    "            if t.creation_op == \"mm\"\n",
    "                c1 = t.creators[1]\n",
    "                c2 = t.creators[2]\n",
    "                new_ =  t.grad * c2' ################\n",
    "                backward(c1, new_)\n",
    "                new_ = c1' * t.grad\n",
    "                backward(c2, new_)\n",
    "            end\n",
    "                  \n",
    "            if t.creation_op == \"transpose\"\n",
    "                backward(t.creators[1], t.grad')\n",
    "            end\n",
    "            \n",
    "            if occursin(\"sum\", t.creation_op)\n",
    "                dim = parse(Int, split(t.creation_op, \"_\")[2])\n",
    "                backward(t.creators[1], expand(t.grad, dim, size(t.creators[1].data)[dim]))\n",
    "            end\n",
    "            \n",
    "            if occursin(\"expand\", t.creation_op)\n",
    "                dim = parse(Int, split(t.creation_op, \"_\")[2])\n",
    "                ndims_cr = ndims(t.creators[1].data)\n",
    "                backward(t.creators[1], dropdims(sum(t.grad;dims=dim);dims=dim, ndims_cr=ndims_cr))\n",
    "            end\n",
    "            \n",
    "            if t.creation_op == \"neg\"\n",
    "                backward(t.creators[1], -t.grad)\n",
    "            end\n",
    "            \n",
    "            if t.creation_op == \"sigmoid\"\n",
    "                ones_ = Tensor_v5(ones(size(t.grad.data)))\n",
    "                backward(t.creators[1], t.grad .* t .* (ones_ - t) )\n",
    "            end\n",
    "            \n",
    "            if t.creation_op == \"tanh\"\n",
    "                ones_ = Tensor_v5(ones(size(t.grad.data)))\n",
    "                backward(t.creators[1], t.grad .* (ones_ - (t .* t)))\n",
    "            end\n",
    "            \n",
    "            if t.creation_op == \"index_select\"\n",
    "                new_grad = zeros(size(t.creators[1]))\n",
    "                indices = t.index_select_indices.data\n",
    "                major_chunks = partition(1:size(t.grad,2),length(indices))\n",
    "                grad_chunks = [t.grad.data[:,inds][:,j]  for(i,inds) in enumerate(major_chunks) for j=1:size(inds)[1]]\n",
    "    \n",
    "                for (i,ind) in enumerate(flatten(indices))\n",
    "                    new_grad[:,ind] +=  grad_chunks[i]\n",
    "                end\n",
    "                backward(t.creators[1], Tensor_v5(new_grad))\n",
    "            end\n",
    "            if t.creation_op == \"cross_entropy\"\n",
    "                dx = t.softmax_output .- t.target_dist\n",
    "                backward(t.creators[1], Tensor_v5(dx))\n",
    "            end\n",
    "        end\n",
    "    end\n",
    "end\n",
    "                        \n",
    "size(a::Tensor_v5) = size(a.data)\n",
    "size(a::Tensor_v5, ind::Int) = size(a.data, ind)\n",
    "\n",
    "function +(a::Tensor_v5, b::Tensor_v5)\n",
    "    if (a.autograd && b.autograd)\n",
    "        return Tensor_v5(a.data + b.data; autograd=true, creators=[a,b], creation_op = \"add\")\n",
    "    end\n",
    "    return Tensor_v5(a.data+b.data)\n",
    "end\n",
    "\n",
    "function -(a::Tensor_v5)\n",
    "    if (a.autograd)\n",
    "        return Tensor_v5(a.data .* -1; autograd=true, creators=[a], creation_op = \"neg\")\n",
    "    end\n",
    "    return Tensor_v5(a.data .* -1)\n",
    "end\n",
    "\n",
    "function -(a::Tensor_v5, b::Tensor_v5)\n",
    "    if (a.autograd && b.autograd)\n",
    "        return Tensor_v5(a.data - b.data; autograd=true, creators=[a,b], creation_op = \"sub\")\n",
    "    end\n",
    "    return Tensor_v5(a.data-b.data)\n",
    "end\n",
    "\n",
    "#element-wise multiplication\n",
    "function broadcasted(f::typeof(*), a::Tensor_v5, b::Tensor_v5)\n",
    "    new_data = zeros(size(a.data))\n",
    "    for i=1:length(new_data)\n",
    "        new_data[i] = f(a.data[i] ,b.data[i])\n",
    "    end\n",
    "    if (a.autograd && b.autograd)\n",
    "        return Tensor_v5(new_data; autograd=true, creators=[a,b], creation_op =\"mul\")\n",
    "    end\n",
    "    return Tensor_v5(new_data)\n",
    "end\n",
    "\n",
    "function broadcasted(f::typeof(-), a::Tensor_v5, b::Tensor_v5)\n",
    "    new_data = zeros(size(a.data))\n",
    "    for i=1:length(new_data)\n",
    "        new_data[i] = -(a.data[i] ,b.data[i])\n",
    "    end\n",
    "    if (a.autograd && b.autograd)\n",
    "        return Tensor_v5(new_data; autograd=true, creators=[a,b], creation_op =\"sub\")\n",
    "    end\n",
    "    return Tensor_v5(new_data)\n",
    "end\n",
    "\n",
    "function sum(a::Tensor_v5; dims=dims)\n",
    "    new_ = dropdims(sum(a.data ;dims=dims), dims = tuple(findall(size(a) .== 1)...))\n",
    "    if (a.autograd)\n",
    "        return Tensor_v5(new_; autograd=true, creators=[a], creation_op = \"sum_\"*string(dims))\n",
    "    end\n",
    "    return Tensor_v5(new_)\n",
    "end\n",
    "\n",
    "function dropdims(a::Tensor_v5;dims=dims,ndims_cr=ndims_cr)\n",
    "    if ndims(a.data) == ndims_cr\n",
    "        return a\n",
    "    end\n",
    "    if (a.autograd)\n",
    "        return Tensor_v5(dropdims(a.data ;dims=dims); autograd=true, creators=[a], creation_op = \"dropdims\")\n",
    "    end\n",
    "    return Tensor_v5(dropdims(a.data ;dims=dims))\n",
    "end\n",
    "\n",
    "function expand(a::Tensor_v5, dim, copies)\n",
    "    sz = size(a)\n",
    "    rep = ntuple(d->d==dim ? copies : 1, length(sz)+1)\n",
    "    new_size = ntuple(d->d<dim ? sz[d] : d == dim ? 1 : sz[d-1], length(sz)+1)\n",
    "    new_data =  repeat(reshape(a.data, new_size), outer=rep)\n",
    "    if (a.autograd)\n",
    "        return Tensor_v5(new_data; autograd=true, creators=[a], creation_op = \"expand_\"*string(dim))\n",
    "    end\n",
    "    return Tensor_v5(new_data)\n",
    "end\n",
    "\n",
    "#transpose\n",
    "function adjoint(a::Tensor_v5)\n",
    "    if (a.autograd)\n",
    "        return Tensor_v5(a.data';autograd=true, creators=[a], creation_op = \"transpose\")\n",
    "    end\n",
    "    return Tensor_v5(a.data')\n",
    "end\n",
    "\n",
    "#matrix multiply \n",
    "function *(a::Tensor_v5, b::Tensor_v5)\n",
    "    if (a.autograd && b.autograd)\n",
    "        return Tensor_v5(a.data * b.data; autograd=true, creators=[a,b], creation_op = \"mm\")\n",
    "    end\n",
    "    return Tensor_v5(a.data * b.data)\n",
    "end\n",
    "\n",
    "\n",
    "function index_select_helper(a::Array, indices)\n",
    "    return reduce(hcat,map(ind -> a[:,ind], indices))\n",
    "end\n",
    "\n",
    "function index_select(a::Tensor_v5, indices::Tensor_v5)\n",
    "    new_ = index_select_helper(a.data, indices.data)\n",
    "    if (a.autograd)\n",
    "        T = Tensor_v5(new_, autograd=true, creators=[a], creation_op = \"index_select\")\n",
    "        T.index_select_indices = indices\n",
    "        return T\n",
    "    end\n",
    "    return Tensor_v5(new_)\n",
    "end\n",
    "\n",
    "println(t::Tensor_v5) = println(t.data)\n",
    "show(io::IO,m::MIME\"text/plain\",a::Tensor_v5) = show(io,m,a.data)\n",
    "                        \n",
    "abstract type Layer end\n",
    "\n",
    "function get_parameters(l::Layer)\n",
    "    return l.parameters\n",
    "end\n",
    "\n",
    "mutable struct Linear <: Layer\n",
    "    W\n",
    "    b\n",
    "    parameters\n",
    "    \n",
    "    function Linear(n_inputs, n_outputs)\n",
    "        W = Tensor_v5(randn(n_outputs, n_inputs) .* sqrt(1.0/n_inputs), autograd=true)\n",
    "        b = Tensor_v5(zeros(n_outputs), autograd=true)\n",
    "        parameters = [W,b]\n",
    "        return new(W,b,parameters)\n",
    "    end\n",
    "end\n",
    "\n",
    "\n",
    "σ(x) = 1/(1+exp(-x))                        \n",
    "\n",
    "struct Tanh <: Layer\n",
    "    Tanh() = new()\n",
    "end\n",
    "\n",
    "struct Sigmoid <: Layer\n",
    "    Sigmoid() = new()\n",
    "end\n",
    "\n",
    "function get_parameters(l::Tanh)\n",
    "    return []\n",
    "end\n",
    "\n",
    "function get_parameters(l::Sigmoid)\n",
    "    return []\n",
    "end\n",
    "\n",
    "function forward(l::Sigmoid, a::Tensor_v5)\n",
    "    if a.autograd\n",
    "        return Tensor_v5(σ.(a.data); autograd=true, creators=[a], creation_op = \"sigmoid\")\n",
    "    end\n",
    "    return Tensor_v5(σ.(a.data))\n",
    "end\n",
    "        \n",
    "function forward(l::Tanh, a::Tensor_v5)\n",
    "    if a.autograd\n",
    "        return Tensor_v5(tanh.(a.data); autograd=true, creators=[a], creation_op = \"tanh\")\n",
    "    end\n",
    "    return Tensor_v5(tanh.(a.data))\n",
    "end    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "mutable struct Embedding_v2 <: Layer\n",
    "    vocab_size\n",
    "    dim\n",
    "    weight\n",
    "    parameters\n",
    "    # this random initialiation style is just a convention from word2vec\n",
    "    function Embedding_v2(dim, vocab_size) \n",
    "        E = new(vocab_size, dim, Tensor_v5((randn(dim, vocab_size) .- 0.5) ./ dim; autograd=true))\n",
    "        E.parameters = [E.weight]\n",
    "        return E\n",
    "    end\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "forward (generic function with 13 methods)"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "using Statistics: mean\n",
    "using LinearAlgebra: I\n",
    "function softmax(x)\n",
    "    temp = exp.(x)\n",
    "    return temp ./ sum(temp;dims=1)\n",
    "end\n",
    "\n",
    "struct CrossEntropyLoss \n",
    "    CrossEntropyLoss() = new()\n",
    "end\n",
    "\n",
    "function forward(l::CrossEntropyLoss, a::Tensor_v5, target::Tensor_v5)\n",
    "    softmax_output = softmax(a.data)\n",
    "    log_out = log.(softmax_output)\n",
    "    sz = size(a.data, 1)\n",
    "    identity = 1.0 .* Matrix(I, (sz, sz))\n",
    "    target_dist = reshape(identity[:,target.data],(size(a.data)))\n",
    "    loss = -mean(sum(log_out .* target_dist;dims=1))\n",
    "    if a.autograd\n",
    "        loss = Tensor_v5(loss; autograd=true, creators=[a], creation_op = \"cross_entropy\")\n",
    "        loss.softmax_output = softmax_output\n",
    "        loss.target_dist = target_dist\n",
    "        return loss\n",
    "    end\n",
    "    return Tensor_v5(loss)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.4461881980473767\n",
      "0.32134282627951916\n",
      "0.07050547863798148\n",
      "0.030112221326762667\n",
      "0.02294860815584766\n",
      "0.01865850502637032\n",
      "0.015762310723879135\n",
      "0.013664053181812933\n",
      "0.012069325190316341\n",
      "0.010814145520457268\n"
     ]
    }
   ],
   "source": [
    "using Random:seed!;seed!(0)\n",
    "data = Tensor_v5([[1,2,1,2]], autograd=true)\n",
    "target = Tensor_v5([4 2 4 2], autograd=true)\n",
    "\n",
    "embed = Embedding_v2(3,3)\n",
    "model = Sequential([embed, Tanh(), Linear(3,4)])\n",
    "criterion = CrossEntropyLoss()\n",
    "\n",
    "\n",
    "optim = SGD(get_parameters(model),1.0)\n",
    "\n",
    "for i=1:10\n",
    "    \n",
    "    # Predict\n",
    "    pred = forward(model, data)\n",
    "    \n",
    "    # Compare\n",
    "    loss = forward(criterion,pred, target)\n",
    "    \n",
    "    # Learn\n",
    "    backward(loss, Tensor_v5(ones(Float32, size(loss.data))))\n",
    "    step(optim)\n",
    "    println(loss)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 17: The Recurrent Neural Network Layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "init_hidden (generic function with 1 method)"
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mutable struct RNNCell <: Layer\n",
    "    n_hidden\n",
    "    \n",
    "    activation\n",
    "    \n",
    "    w_ih\n",
    "    w_hh\n",
    "    w_ho\n",
    "    \n",
    "    parameters\n",
    "    \n",
    "    function RNNCell(n_inputs, n_hidden, n_output, activation=\"sigmoid\")\n",
    "        if activation == \"sigmoid\"\n",
    "            act = Sigmoid()\n",
    "        elseif activation == \"tanh\"\n",
    "            act = Tanh()\n",
    "        else\n",
    "            throw(\"Non-linearity not found\")\n",
    "        end\n",
    "        \n",
    "        parameters = []\n",
    "\n",
    "        w_ih = Linear(n_inputs, n_hidden)\n",
    "        w_hh = Linear(n_hidden, n_hidden)\n",
    "        w_ho = Linear(n_hidden, n_output)\n",
    "        \n",
    "        push!(parameters, get_parameters(w_ih))\n",
    "        push!(parameters, get_parameters(w_hh))\n",
    "        push!(parameters, get_parameters(w_ho))\n",
    "        parameters = collect(Iterators.flatten(parameters))\n",
    "        return new(n_hidden, act, w_ih, w_hh, w_ho, parameters)\n",
    "    end\n",
    "end\n",
    "\n",
    "function forward(rnn::RNNCell, input::Tensor_v5, hidden::Tensor_v5)\n",
    "    from_prev_hidden = forward(rnn.w_hh, hidden)\n",
    "    combined = forward(rnn.w_ih, input) + from_prev_hidden\n",
    "    new_hidden = forward(rnn.activation, combined)\n",
    "    output = forward(rnn.w_ho, new_hidden)\n",
    "    return output, new_hidden\n",
    "end\n",
    "\n",
    "function init_hidden(rnn::RNNCell; batch_size=1)\n",
    "    return Tensor_v5(zeros(rnn.n_hidden, batch_size), autograd=true)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "raw = readlines(\"tasksv11/en/qa1_single-supporting-fact_train.txt\")\n",
    "\n",
    "tokens = []\n",
    "for line in raw[1:1000]\n",
    "    push!(tokens, split(lowercase(line),\" \")[2:end])\n",
    "end\n",
    "\n",
    "new_tokens = []\n",
    "for line in tokens\n",
    "    push!(new_tokens, cat(repeat([\"-\"],6-length(line)), line;dims=1))\n",
    "end\n",
    "\n",
    "tokens = new_tokens\n",
    "\n",
    "vocab = Set()\n",
    "for sent in tokens\n",
    "    for word in sent\n",
    "        if length(word)>0\n",
    "            push!(vocab, word)\n",
    "        end\n",
    "    end\n",
    "end\n",
    "vocab = collect(vocab)\n",
    "\n",
    "word2index = Dict()\n",
    "for (i,word) in enumerate(vocab)\n",
    "    word2index[word] = i\n",
    "end\n",
    "\n",
    "indices = []\n",
    "for line in tokens\n",
    "    idx = []\n",
    "    for w in line\n",
    "        push!(idx,word2index[w])\n",
    "    end\n",
    "    push!(indices, idx)\n",
    "end\n",
    "\n",
    "data = reduce(hcat,indices);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "embed = Embedding_v2(16, length(vocab))\n",
    "model = RNNCell(16, 16, length(vocab))\n",
    "criterion = CrossEntropyLoss()\n",
    "optim = SGD(cat(get_parameters(model), get_parameters(embed); dims=1), 0.05);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loss: 0.44647163118068134 correct: 0.0\n",
      "Loss: 0.18065613931415564 correct: 0.21\n",
      "Loss: 0.17841670806583004 correct: 0.22\n",
      "Loss: 0.16513976605062758 correct: 0.31\n",
      "Loss: 0.1459613675665265 correct: 0.36\n"
     ]
    }
   ],
   "source": [
    "for iter=1:1000\n",
    "    batch_size = 100\n",
    "    total_loss = 0\n",
    "    \n",
    "    hidden = init_hidden(model, batch_size=batch_size)\n",
    "    output = nothing #to access from for loop\n",
    "    for t=1:5\n",
    "        input = Tensor_v5(data[t,1:batch_size], autograd=true)\n",
    "        rnn_input = forward(embed, input)\n",
    "        output, hidden = forward(model, rnn_input, hidden)\n",
    "    end\n",
    "    target = Tensor_v5(data[6,1:batch_size], autograd=true)\n",
    "    loss = forward(criterion, output, target)\n",
    "    backward(loss)\n",
    "    step(optim)\n",
    "    total_loss += loss.data\n",
    "    \n",
    "    if (iter-1) %200 ==0\n",
    "        max_ind = argmax(output.data;dims=1)\n",
    "        p_correct = dropdims(map(x->x.I[1], max_ind);dims=1)\n",
    "        p_correct = mean(target.data .== p_correct)\n",
    "        println(\"Loss: $(total_loss/10) correct: $(p_correct)\")\n",
    "    end\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Context: - mary moved to the \n",
      "True: bathroom.\n",
      "Pred: garden.\n"
     ]
    }
   ],
   "source": [
    "batch_size = 1\n",
    "hidden = init_hidden(model, batch_size=batch_size)\n",
    "\n",
    "output = nothing #to access from for loop\n",
    "for t=1:5\n",
    "    input = Tensor_v5(data[t,1:batch_size], autograd=true)\n",
    "    rnn_input = forward(embed, input)\n",
    "    output, hidden = forward(model, rnn_input, hidden)\n",
    "end\n",
    "\n",
    "target = Tensor_v5(data[6,1:batch_size], autograd=true)   \n",
    "loss = forward(criterion, output, target)\n",
    "\n",
    "ctx = \"\"\n",
    "for idx in data[:,1][1:end-1]\n",
    "    global ctx *= vocab[idx] * \" \"\n",
    "end\n",
    "println(\"Context: \",ctx)\n",
    "println(\"True: \",vocab[target.data[1]])\n",
    "println(\"Pred: \", vocab[argmax(output.data).I[1]])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.4.0",
   "language": "julia",
   "name": "julia-1.4"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.4.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}