notebooks/stacks.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Stacks\n",
    "\n",
    "There exist many implementations of differentiable stacks in the literature related to neural turing machines and similar. e.g. [Learning to Transduce with Unbounded Memory](http://papers.nips.cc/paper/5648-learning-to-transduce-with-unbounded-memory.pdf), [Inferring Algorithmic Patterns withStack-Augmented Recurrent Nets](https://papers.nips.cc/paper/5857-inferring-algorithmic-patterns-with-stack-augmented-recurrent-nets.pdf) e.t.c.\n",
    "\n",
    "However, here we restrict ourselves to follow these two rules. \n",
    "1. It must be deterministic and lossless in forward pass\n",
    "2. It must have well definied gradients in backward pass\n",
    "\n",
    "Tensorflow's autograd does a good job at keeping track where the gradients should flow and thus the differentiable implementation looks almost identical to the classical implementation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stack representation\n",
    "\n",
    "The `stack` variable has two variables, buffer and index. The buffer is the writable buffer where stack elements are stored. Index points to top of stack + 1.\n",
    "\n",
    "Note: Since these functions can create variables, they must execute eagerly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(<tf.Tensor: shape=(3, 3), dtype=float32, numpy=\n",
      "array([[0., 0., 0.],\n",
      "       [0., 0., 0.],\n",
      "       [0., 0., 0.]], dtype=float32)>, <tf.Tensor: shape=(3,), dtype=float32, numpy=array([1., 0., 0.], dtype=float32)>)\n",
      "(<tf.Variable 'Variable:0' shape=(3, 3, 3) dtype=float32, numpy=\n",
      "array([[[0., 0., 0.],\n",
      "        [0., 0., 0.],\n",
      "        [0., 0., 0.]],\n",
      "\n",
      "       [[0., 0., 0.],\n",
      "        [0., 0., 0.],\n",
      "        [0., 0., 0.]],\n",
      "\n",
      "       [[0., 0., 0.],\n",
      "        [0., 0., 0.],\n",
      "        [0., 0., 0.]]], dtype=float32)>, <tf.Variable 'Variable:0' shape=(3,) dtype=float32, numpy=array([1., 0., 0.], dtype=float32)>)\n"
     ]
    }
   ],
   "source": [
    "def new_stack(stack_shape, is_learnable=False):\n",
    "    buffer = tf.zeros(stack_shape, dtype=tf.float32)\n",
    "    index = tf.one_hot(0, stack_shape[0], dtype=tf.float32)\n",
    "    \n",
    "    if is_learnable:\n",
    "        buffer = tf.Variable(buffer)\n",
    "        index = tf.Variable(index)\n",
    "    \n",
    "    stack = (buffer, index)\n",
    "    return stack\n",
    "\n",
    "constant_stack = new_stack((3,3))\n",
    "print(constant_stack)\n",
    "\n",
    "learnable_stack = new_stack((3,3,3), True)\n",
    "print(learnable_stack)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(<tf.Variable 'Variable:0' shape=(3, 3) dtype=float32, numpy=\n",
       " array([[1., 1., 1.],\n",
       "        [1., 1., 1.],\n",
       "        [1., 1., 1.]], dtype=float32)>,\n",
       " <tf.Variable 'Variable:0' shape=(3,) dtype=float32, numpy=array([1., 0., 0.], dtype=float32)>)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def new_stack_from_buffer(buffer, is_learnable=False):\n",
    "    stack_shape = tf.shape(buffer)\n",
    "    index = tf.one_hot(0, stack_shape[0], dtype=tf.float32)\n",
    "    \n",
    "    if is_learnable:\n",
    "        buffer = tf.Variable(buffer)\n",
    "        index = tf.Variable(index)\n",
    "    \n",
    "    stack = (buffer, index)\n",
    "    return stack\n",
    "\n",
    "buffer = tf.ones((3,3), dtype=tf.float32)\n",
    "stack = new_stack_from_buffer(buffer, True)\n",
    "stack"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Soft assignment\n",
    "Tensorflow does not allow direct assignment of array indexes, so we use this trick. For more information go to [bubble-sort.ipynb](bubble-sort.ipynb) or [differentiable-indexed-arrays.ipynb](differentiable-indexed-arrays.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from library.array_ops import assign_index_vectored, superposition_lookup_vectored"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stack push\n",
    "The `stack_push` function is a stateless function. At the time of writing, the Autograph has undefined behaviour if we try to build a stateful implementation of stack like using python class or using closures."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[[1. 1. 1.]\n",
      " [2. 2. 2.]\n",
      " [3. 3. 3.]], shape=(3, 3), dtype=float32)\n",
      "tf.Tensor([1. 0. 0.], shape=(3,), dtype=float32)\n",
      "tf.Tensor(\n",
      "[[1. 1. 1.]\n",
      " [1. 1. 1.]\n",
      " [1. 1. 1.]], shape=(3, 3), dtype=float32)\n",
      "(None, None)\n"
     ]
    }
   ],
   "source": [
    "@tf.function\n",
    "def stack_push(stack, element):\n",
    "    buffer, index = stack\n",
    "    buffer = assign_index_vectored(buffer, index, element)\n",
    "    index = tf.roll(index, shift=1, axis=0)\n",
    "    stack = (buffer, index)\n",
    "    return stack\n",
    "\n",
    "stack = new_stack((3,3))\n",
    "elements = tf.Variable([\n",
    "    [1,1,1],\n",
    "    [2,2,2],\n",
    "    [3,3,3]\n",
    "],dtype=tf.float32)\n",
    "\n",
    "original_stack = stack\n",
    "\n",
    "with tf.GradientTape(persistent=True) as tape:\n",
    "    stack = stack_push(stack, elements[0])\n",
    "    stack = stack_push(stack, elements[1])\n",
    "    stack = stack_push(stack, elements[2])\n",
    "    \n",
    "print(stack[0])\n",
    "print(stack[1])\n",
    "print(tape.gradient(stack[0], elements))\n",
    "print(tape.gradient(stack, original_stack))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stack pop\n",
    "For buffer lookup we use the `superposition_lookup_vectored` as described in [differentiable-indexed-arrays.ipynb](differentiable-indexed-arrays.ipynb). We also update the index and return both state and element."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor([3. 3. 3.], shape=(3,), dtype=float32)\n",
      "tf.Tensor([2. 2. 2.], shape=(3,), dtype=float32)\n",
      "tf.Tensor(\n",
      "[[1. 1. 1.]\n",
      " [2. 2. 2.]\n",
      " [3. 3. 3.]], shape=(3, 3), dtype=float32)\n",
      "tf.Tensor([0. 1. 0.], shape=(3,), dtype=float32)\n",
      "None\n"
     ]
    }
   ],
   "source": [
    "@tf.function\n",
    "def stack_pop(stack):\n",
    "    buffer, index = stack\n",
    "    index = tf.roll(index, shift=-1, axis=0)\n",
    "    element = superposition_lookup_vectored(buffer, index)\n",
    "    stack = (buffer, index)\n",
    "    return stack, element\n",
    "\n",
    "buffer = tf.Variable([\n",
    "    [1,1,1],\n",
    "    [2,2,2],\n",
    "    [3,3,3]\n",
    "],dtype=tf.float32)\n",
    "stack = new_stack_from_buffer(buffer, True)\n",
    "\n",
    "with tf.GradientTape() as tape:\n",
    "    ns1, element = stack_pop(stack)\n",
    "    print(element)\n",
    "    ns2, element = stack_pop(ns1)\n",
    "    print(element)\n",
    "\n",
    "print(ns2[0])\n",
    "print(ns2[1])\n",
    "print(tape.gradient(element, buffer))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stack peek\n",
    "\n",
    "Get the stack top without any modification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor([3. 3. 3.], shape=(3,), dtype=float32)\n",
      "None\n"
     ]
    }
   ],
   "source": [
    "@tf.function\n",
    "def stack_peek(stack):\n",
    "    buffer, index = stack\n",
    "    index = tf.roll(index, shift=-1, axis=0)\n",
    "    element = superposition_lookup_vectored(buffer, index)\n",
    "    return element\n",
    "\n",
    "buffer = tf.Variable([\n",
    "    [1,1,1],\n",
    "    [2,2,2],\n",
    "    [3,3,3]\n",
    "],dtype=tf.float32)\n",
    "stack = new_stack_from_buffer(buffer, True)\n",
    "\n",
    "with tf.GradientTape() as tape:\n",
    "    element = stack_peek(stack)\n",
    "\n",
    "print(element)\n",
    "print(tape.gradient(element, buffer))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Toy example: Reversing a list\n",
    "Using two stacks, we can reverse a list. The algorithm has two steps\n",
    "* Stack 1 pushes all elements into itself\n",
    "* Stack 1 then pops an element and Stack 2 pushes that element into itself\n",
    "\n",
    "The buffer of Stack 2 is the solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[[4. 4. 4. 4.]\n",
      " [3. 3. 3. 3.]\n",
      " [2. 2. 2. 2.]\n",
      " [1. 1. 1. 1.]], shape=(4, 4), dtype=float32)\n",
      "tf.Tensor(\n",
      "[[1. 1. 1. 1.]\n",
      " [1. 1. 1. 1.]\n",
      " [1. 1. 1. 1.]\n",
      " [1. 1. 1. 1.]], shape=(4, 4), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "@tf.function\n",
    "def reverse_list(arr):\n",
    "    arr_shape = tf.shape(arr)\n",
    "    arr = tf.unstack(arr)\n",
    "    \n",
    "    stack1 = new_stack(arr_shape)\n",
    "    \n",
    "    # Step 1: Push all elements into stack 1\n",
    "    for element in arr:\n",
    "        stack1 = stack_push(stack1, element)\n",
    "    \n",
    "    stack2 = new_stack(arr_shape)\n",
    "    \n",
    "    # Step 2: Transfer all elements to stack 2\n",
    "    for _ in tf.range(arr_shape[0]):\n",
    "        stack1, element = stack_pop(stack1)\n",
    "        stack2 = stack_push(stack2, element)\n",
    "    \n",
    "    # Return buffer of stack 2\n",
    "    return stack2[0]\n",
    "\n",
    "arr = tf.Variable([\n",
    "    [1,1,1,1],\n",
    "    [2,2,2,2],\n",
    "    [3,3,3,3],\n",
    "    [4,4,4,4],\n",
    "], dtype=tf.float32)\n",
    "\n",
    "with tf.GradientTape() as tape:\n",
    "    new_arr = reverse_list(arr)\n",
    "\n",
    "print(new_arr)\n",
    "print(tape.gradient(new_arr, arr))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Backward pass\n",
    "To demonstrate the working of backward pass, we give a reversed target array `reversed_arr` to the algorithm and a learnable `input_arr`. The algorithm must learn the `input_arr` using gradients."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "28\n",
      "10.2250357\n",
      "2.75192738\n",
      "0.439089298\n",
      "0.129742727\n",
      "0.0685746\n",
      "0.0525279418\n",
      "0.0155251706\n",
      "0.000200064795\n",
      "0.00172046362\n",
      "[[1 1 1 1]\n",
      " [2 2 2 2]\n",
      " [3 3 3 3]\n",
      " [4 4 4 4]]\n"
     ]
    }
   ],
   "source": [
    "opt = tf.keras.optimizers.Adam(1e-1)\n",
    "\n",
    "@tf.function\n",
    "def train_step(x, y):\n",
    "    with tf.GradientTape() as tape:\n",
    "        y_ = reverse_list(x)\n",
    "        loss = tf.nn.l2_loss(y - y_)\n",
    "        \n",
    "    grads = tape.gradient(loss, x)\n",
    "    opt.apply_gradients(zip([grads], [x]))\n",
    "    \n",
    "    return loss\n",
    "\n",
    "input_arr = tf.Variable([\n",
    "    [1,1,1,1],\n",
    "    [1,1,1,1],\n",
    "    [1,1,1,1],\n",
    "    [1,1,1,1]\n",
    "], dtype=tf.float32)\n",
    "reversed_arr = tf.constant([\n",
    "    [4,4,4,4],\n",
    "    [3,3,3,3],\n",
    "    [2,2,2,2],\n",
    "    [1,1,1,1],\n",
    "], dtype=tf.float32)\n",
    "\n",
    "for i in range(100):\n",
    "    loss = train_step(input_arr, reversed_arr)\n",
    "    if i % 10 == 0:\n",
    "        tf.print(loss)\n",
    "tf.print(tf.round(input_arr))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7.6 64-bit",
   "language": "python",
   "name": "python37664bit238183a89ccd4c25acc508071275f29e"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}