{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction to Scan in Theano\n",
    "\n",
    "Credits: Forked from [summerschool2015](https://github.com/mila-udem/summerschool2015) by mila-udem\n",
    "\n",
    "## In short\n",
    "\n",
    "* Mechanism to perform loops in a Theano graph\n",
    "* Supports nested loops and reusing results from previous iterations \n",
    "* Highly generic\n",
    "\n",
    "## Implementation\n",
    "\n",
    "A Theano function graph is composed of two types of nodes; Variable nodes which represent data and Apply node which apply Ops (which represent some computation) to Variables to produce new Variables.\n",
    "\n",
    "From this point of view, a node that applies a Scan op is just like any other. Internally, however, it is very different from most Ops.\n",
    "\n",
    "Inside a Scan op is yet another Theano graph which represents the computation to be performed at every iteration of the loop. During compilation, that graph is compiled into a function and, during execution, the Scan op will call that function repeatedly on its inputs to produce its outputs.\n",
    "\n",
    "## Example 1 : As simple as it gets\n",
    "\n",
    "Scan's interface is complex and, thus, best introduced by examples. So, let's dive right in and start with a simple example; perform an element-wise multiplication between two vectors. \n",
    "\n",
    "This particular example is simple enough that Scan is not the best way to do things but we'll gradually work our way to more complex examples where Scan gets more interesting.\n",
    "\n",
    "Let's first setup our use case by defining Theano variables for the inputs :"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import theano\n",
    "import theano.tensor as T\n",
    "import numpy as np\n",
    "\n",
    "vector1 = T.vector('vector1')\n",
    "vector2 = T.vector('vector2')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we call the `scan()` function. It has many parameters but, because our use case is simple, we only need two of them. We'll introduce other parameters in the next examples.\n",
    "\n",
    "The parameter `sequences` allows us to specify variables that Scan should iterate over as it loops. The first iteration will take as input the first element of every sequence, the second iteration will take as input the second element of every sequence, etc. These individual element have will have one less dimension than the original sequences. For example, for a matrix sequence, the individual elements will be vectors.\n",
    "\n",
    "The parameter `fn` receives a function or lambda expression that expresses the computation to do at every iteration. It operates on the symbolic inputs to produce symbolic outputs. It will **only ever be called once**, to assemble the Theano graph used by Scan at every the iterations.\n",
    "\n",
    "Since we wish to iterate over both `vector1` and `vector2` simultaneously, we provide them as sequences. This means that every iteration will operate on two inputs: an element from `vector1` and the corresponding element from `vector2`. \n",
    "\n",
    "Because what we want is the elementwise product between the vectors, we provide a lambda expression that, given an element `a` from `vector1` and an element `b` from `vector2` computes and return the product."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "output, updates = theano.scan(fn=lambda a, b : a * b,\n",
    "                              sequences=[vector1, vector2])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calling `scan()`, we see that it returns two outputs.\n",
    "\n",
    "The first output contains the outputs of `fn` from every timestep concatenated into a tensor. In our case, the output of a single timestep is a scalar so output is a vector where `output[i]` is the output of the i-th iteration.\n",
    "\n",
    "The second output details if and how the execution of the Scan updates any shared variable in the graph. It should be provided as an argument when compiling the Theano function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "f = theano.function(inputs=[vector1, vector2],\n",
    "                    outputs=output,\n",
    "                    updates=updates)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If `updates` is omitted, the state of any shared variables modified by Scan will not be updated properly. Random number sampling, for instance, relies on shared variables. If `updates` is not provided, the state of the random number generator won't be updated properly and the same numbers might be sampled repeatedly. **Always** provide `updates` when compiling your Theano function.\n",
    "\n",
    "Now that we've defined how to do elementwise multiplication with Scan, we can see that the result is as expected :"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "vector1_value = np.arange(0, 5).astype(theano.config.floatX) # [0,1,2,3,4]\n",
    "vector2_value = np.arange(1, 6).astype(theano.config.floatX) # [1,2,3,4,5]\n",
    "print(f(vector1_value, vector2_value))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "An interesting thing is that we never explicitly told Scan how many iteration it needed to run. It was automatically inferred; when given sequences, Scan will run as many iterations as the length of the shortest sequence : "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print(f(vector1_value, vector2_value[:4]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example 2 : Non-sequences\n",
    "\n",
    "In this example, we introduce another of Scan's features; non-sequences. To demonstrate how to use them, we use Scan to compute the activations of  a linear MLP layer over a minibatch.\n",
    "\n",
    "It is not yet a use case where Scan is truly useful but it introduces a requirement that sequences cannot fulfill; if we want to use Scan to iterate over the minibatch elements and compute the activations for each of them, then we need some variables (the parameters of the layer), to be available 'as is' at every iteration of the loop. We do *not* want Scan to iterate over them and give only part of them at every iteration.\n",
    "\n",
    "Once again, we begin by setting up our Theano variables :"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "X = T.matrix('X') # Minibatch of data\n",
    "W = T.matrix('W') # Weights of the layer\n",
    "b = T.vector('b') # Biases of the layer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the sake of variety, in this example we define the computation to be done at every iteration of the loop using a Python function, `step()`, instead of a lambda expression.\n",
    "\n",
    "To have the full weight matrix W and the full bias vector b available at every iteration, we use the argument non_sequences. Contrary to sequences, non-sequences are not iterated upon by Scan. Every non-sequence is passed as input to every iteration.\n",
    "\n",
    "This means that our `step()` function will need to operate on three symbolic inputs; one for our sequence X and one for each of our non-sequences W and b. \n",
    "\n",
    "The inputs that correspond to the non-sequences are **always** last and in the same order at the non-sequences are provided to Scan. This means that the correspondence between the inputs of the `step()` function and the arguments to `scan()` is the following : \n",
    "\n",
    "* `v` : individual element of the sequence `X` \n",
    "* `W` and `b` : non-sequences `W` and `b`, respectively"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def step(v, W, b):\n",
    "    return T.dot(v, W) + b\n",
    "\n",
    "output, updates = theano.scan(fn=step,\n",
    "                              sequences=[X],\n",
    "                              non_sequences=[W, b])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now compile our Theano function and see that it gives the expected results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "f = theano.function(inputs=[X, W, b],\n",
    "                    outputs=output,\n",
    "                    updates=updates)\n",
    "\n",
    "X_value = np.arange(-3, 3).reshape(3, 2).astype(theano.config.floatX)\n",
    "W_value = np.eye(2).astype(theano.config.floatX)\n",
    "b_value = np.arange(2).astype(theano.config.floatX)\n",
    "print(f(X_value, W_value, b_value))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example 3 : Reusing outputs from the previous iterations\n",
    "\n",
    "In this example, we will use Scan to compute a cumulative sum over the first dimension of a matrix $M$. This means that the output will be a matrix $S$ in which the first row will be equal to the first row of $M$, the second row will be equal to the sum of the two first rows of $M$, and so on.\n",
    "\n",
    "Another way to express this, which is the way we will implement here, is that $S[t] = S[t-1] + M[t]$. Implementing this with Scan would involve iterating over the rows of the matrix $M$ and, at every iteration, reuse the cumulative row that was output at the previous iteration and return the sum of it and the current row of $M$.\n",
    "\n",
    "If we assume for a moment that we can get Scan to provide the output value from the previous iteration as an input for every iteration, implementing a step function is simple :"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def step(m_row, cumulative_sum):\n",
    "    return m_row + cumulative_sum"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The trick part is informing Scan that our step function expects as input the output of a previous iteration. To achieve this, we need to use a new parameter of the `scan()` function: `outputs_info`. This parameter is used to tell Scan how we intend to use each of the outputs that are computed at each iteration.\n",
    "\n",
    "This parameter can be omitted (like we did so far) when the step function doesn't depend on any output of a previous iteration. However, now that we wish to have recurrent outputs, we need to start using it.\n",
    "\n",
    "`outputs_info` takes a sequence with one element for every output of the `step()` function :\n",
    "* For a **non-recurrent output** (like in every example before this one), the element should be `None`.\n",
    "* For a **simple recurrent output** (iteration $t$ depends on the value at iteration $t-1$), the element must be a tensor. Scan will interpret it as being an initial state for a recurrent output and give it as input to the first iteration, pretending it is the output value from a previous iteration. For subsequent iterations, Scan will automatically handle giving the previous output value as an input.\n",
    "\n",
    "The `step()` function needs to expect one additional input for each simple recurrent output. These inputs correspond to outputs from previous iteration and are **always** after the inputs that correspond to sequences but before those that correspond to non-sequences. The are received by the `step()` function in the order in which the recurrent outputs are declared in the outputs_info sequence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "M = T.matrix('X')\n",
    "s = T.vector('s') # Initial value for the cumulative sum\n",
    "\n",
    "output, updates = theano.scan(fn=step,\n",
    "                              sequences=[M],\n",
    "                              outputs_info=[s])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now compile and test the Theano function :"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "f = theano.function(inputs=[M, s],\n",
    "                    outputs=output,\n",
    "                    updates=updates)\n",
    "\n",
    "M_value = np.arange(9).reshape(3, 3).astype(theano.config.floatX)\n",
    "s_value = np.zeros((3, ), dtype=theano.config.floatX)\n",
    "print(f(M_value, s_value))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "An important thing to notice here, is that the output computed by the Scan does **not** include the initial state that we provided. It only outputs the states that it has computed itself.\n",
    "\n",
    "If we want to have both the initial state and the computed states in the same Theano variable, we have to join them ourselves."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example 4 : Reusing outputs from multiple past iterations\n",
    "\n",
    "The Fibonacci sequence is a sequence of numbers F where the two first numbers both 1 and every subsequence number is defined as such : $F_n = F_{n-1} + F_{n-2}$. Thus, the Fibonacci sequence goes : 1, 1, 2, 3, 5, 8, 13, ...\n",
    "\n",
    "In this example, we will cover how to compute part of the Fibonacci sequence using Scan. Most of the tools required to achieve this have been introduced in the previous examples. The only one missing is the ability to use, at iteration $i$, outputs from iterations older than $i-1$.\n",
    "\n",
    "Also, since every example so far had only one output at every iteration of the loop, we will also compute, at each timestep, the ratio between the new term of the Fibonacci sequence and the previous term.\n",
    "\n",
    "Writing an appropriate step function given two inputs, representing the two previous terms of the Fibonacci sequence, is easy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def step(f_minus2, f_minus1):\n",
    "    new_f = f_minus2 + f_minus1\n",
    "    ratio = new_f / f_minus1\n",
    "    return new_f, ratio"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next step is defining the value of `outputs_info`.\n",
    "\n",
    "Recall that, for **non-recurrent outputs**, the value is `None` and, for **simple recurrent outputs**, the value is a single initial state. For **general recurrent outputs**, where iteration $t$ may depend on multiple past values, the value is a dictionary. That dictionary has two values:\n",
    "* taps : list declaring which previous values of that output every iteration will need. `[-3, -2, -1]` would mean every iteration should take as input the last 3 values of that output. `[-2]` would mean every iteration should take as input the value of that output from two iterations ago.\n",
    "* initial : tensor of initial values. If every initial value has $n$ dimensions, `initial` will be a single tensor of $n+1$ dimensions with as many initial values as the oldest requested tap. In the case of the Fibonacci sequence, the individual initial values are scalars so the `initial` will be a vector. \n",
    "\n",
    "In our example, we have two outputs. The first output is the next computed term of the Fibonacci sequence so every iteration should take as input the two last values of that output. The second output is the ratio between successive terms and we don't reuse its value so this output is non-recurrent. We define the value of `outputs_info` as such :"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "f_init = T.fvector()\n",
    "outputs_info = [dict(initial=f_init, taps=[-2, -1]),\n",
    "                None]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we've defined the step function and the properties of our outputs, we can call the `scan()` function. Because the `step()` function has multiple outputs, the first output of `scan()` function will be a list of tensors: the first tensor containing all the states of the first output and the second tensor containing all the states of the second input.\n",
    "\n",
    "In every previous example, we used sequences and Scan automatically inferred the number of iterations it needed to run from the length of these\n",
    "sequences. Now that we have no sequence, we need to explicitly tell Scan how many iterations to run using the `n_step` parameter. The value can be real or symbolic."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "output, updates = theano.scan(fn=step,\n",
    "                              outputs_info=outputs_info,\n",
    "                              n_steps=10)\n",
    "\n",
    "next_fibonacci_terms = output[0]\n",
    "ratios_between_terms = output[1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's compile our Theano function which will take a vector of consecutive values from the Fibonacci sequence and compute the next 10 values :"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "f = theano.function(inputs=[f_init],\n",
    "                    outputs=[next_fibonacci_terms, ratios_between_terms],\n",
    "                    updates=updates)\n",
    "\n",
    "out = f([1, 1])\n",
    "print(out[0])\n",
    "print(out[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Precisions about the order of the arguments to the step function\n",
    "\n",
    "When we start using many sequences, recurrent outputs and non-sequences, it's easy to get confused regarding the order in which the step function receives the corresponding inputs. Below is the full order:\n",
    "\n",
    "* Element from the first sequence\n",
    "* ...\n",
    "* Element from the last sequence\n",
    "* First requested tap from first recurrent output\n",
    "* ...\n",
    "* Last requested tap from first recurrent output\n",
    "* ...\n",
    "* First requested tap from last recurrent output\n",
    "* ...\n",
    "* Last requested tap from last recurrent output\n",
    "* First non-sequence\n",
    "* ...\n",
    "* Last non-sequence"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## When to use Scan and when not to\n",
    "\n",
    "Scan is not appropriate for every problem. Here's some information to help you figure out if Scan is the best solution for a given use case.\n",
    "\n",
    "### Execution speed\n",
    "\n",
    "Using Scan in a Theano function typically makes it slighly slower compared to the equivalent Theano graph in which the loop is unrolled. Both of these approaches tend to be much slower than a vectorized implementation in which large chunks of the computation can be done in parallel.\n",
    "\n",
    "### Compilation speed\n",
    "\n",
    "Scan also adds an overhead to the compilation, potentially making it slower, but using it can also dramatically reduce the size of your graph, making compilation much faster. In the end, the effect of Scan on compilation speed will heavily depend on the size of the graph with and without Scan.\n",
    "\n",
    "The compilation speed of a Theano function using Scan will usually be comparable to one in which the loop is unrolled if the number of iterations is small. It the number of iterations is large, however, the compilation will usually be much faster with Scan.\n",
    "\n",
    "### In summary\n",
    "\n",
    "If you have one of the following cases, Scan can help :\n",
    "* A vectorized implementation is not possible (due to the nature of the computation and/or memory usage)\n",
    "* You want to do a large or variable number of iterations\n",
    "\n",
    "If you have one of the following cases, you should consider other options :\n",
    "* A vectorized implementation could perform the same computation => Use the vectorized approach. It will often be faster during both compilation and execution.\n",
    "* You want to do a small, fixed, number of iterations (ex: 2 or 3) => It's probably better to simply unroll the computation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercises\n",
    "\n",
    "### Exercise 1 - Computing a polynomial\n",
    "\n",
    "In this exercise, the initial version already works. It computes the value of a polynomial ($n_0 + n_1 x + n_2 x^2 + ... $) of at most 10000 degrees given the coefficients of the various terms and the value of x.\n",
    "\n",
    "You must modify it such that the reduction (the sum() call) is done by Scan."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "coefficients = theano.tensor.vector(\"coefficients\")\n",
    "x = T.scalar(\"x\")\n",
    "max_coefficients_supported = 10000\n",
    "\n",
    "def step(coeff, power, free_var):\n",
    "    return coeff * free_var ** power\n",
    "\n",
    "# Generate the components of the polynomial\n",
    "full_range=theano.tensor.arange(max_coefficients_supported)\n",
    "components, updates = theano.scan(fn=step,\n",
    "                                  outputs_info=None,\n",
    "                                  sequences=[coefficients, full_range],\n",
    "                                  non_sequences=x)\n",
    "\n",
    "polynomial = components.sum()\n",
    "calculate_polynomial = theano.function(inputs=[coefficients, x],\n",
    "                                       outputs=polynomial,\n",
    "                                       updates=updates)\n",
    "\n",
    "test_coeff = np.asarray([1, 0, 2], dtype=theano.config.floatX)\n",
    "print(calculate_polynomial(test_coeff, 3))\n",
    "# 19.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Solution** : run the cell below to display the solution to this exercise."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%load scan_ex1_solution.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 2 - Sampling without replacement\n",
    "\n",
    "In this exercise, the goal is to implement a Theano function that :\n",
    "* takes as input a vector of probabilities and a scalar\n",
    "* performs sampling without replacements from those probabilities as many times as the value of the scalar\n",
    "* returns a vector containing the indices of the sampled elements.\n",
    "\n",
    "Partial code is provided to help with the sampling of random numbers since this is not something that was covered in this tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "probabilities = T.vector()\n",
    "nb_samples = T.iscalar()\n",
    "\n",
    "rng = T.shared_randomstreams.RandomStreams(1234)\n",
    "\n",
    "def sample_from_pvect(pvect):\n",
    "    \"\"\" Provided utility function: given a symbolic vector of\n",
    "    probabilities (which MUST sum to 1), sample one element\n",
    "    and return its index.\n",
    "    \"\"\"\n",
    "    onehot_sample = rng.multinomial(n=1, pvals=pvect)\n",
    "    sample = onehot_sample.argmax()\n",
    "    return sample\n",
    "\n",
    "def set_p_to_zero(pvect, i):\n",
    "    \"\"\" Provided utility function: given a symbolic vector of\n",
    "    probabilities and an index 'i', set the probability of the\n",
    "    i-th element to 0 and renormalize the probabilities so they\n",
    "    sum to 1.\n",
    "    \"\"\"\n",
    "    new_pvect = T.set_subtensor(pvect[i], 0.)\n",
    "    new_pvect = new_pvect / new_pvect.sum()\n",
    "    return new_pvect\n",
    "    \n",
    "\n",
    "# TODO use Scan to sample from the vector of probabilities and\n",
    "# symbolically obtain 'samples' the vector of sampled indices.\n",
    "samples = None\n",
    "\n",
    "# Compiling the function\n",
    "f = theano.function(inputs=[probabilities, nb_samples],\n",
    "                    outputs=[samples])\n",
    "\n",
    "# Testing the function\n",
    "test_probs = np.asarray([0.6, 0.3, 0.1], dtype=theano.config.floatX)\n",
    "for i in range(10):\n",
    "    print(f(test_probs, 2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Solution** : run the cell below to display the solution to this exercise."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%load scan_ex2_solution.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.4.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}