Manual Theano stack.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import theano\n",
    "from theano import tensor as T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "num_steps = 5\n",
    "batch_size = 3\n",
    "emb_dim = 7\n",
    "\n",
    "batch_range = T.arange(batch_size, dtype=\"int32\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "stack = theano.shared(np.zeros((num_steps * batch_size, emb_dim), dtype=np.float32), name=\"stack\")\n",
    "stack_grad = theano.shared(np.zeros_like(stack.get_value()), name=\"stack_grad\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "cursors = T.ones((batch_size,), dtype=np.int32) * -1\n",
    "queue = T.zeros((batch_size * num_steps,), dtype=np.int32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "Xs = T.tensor3(\"Xs\")\n",
    "masks = T.imatrix(\"masks\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def stack_op(t, stack_t, stack_grad_t, push_value, merge_value,\n",
    "             queue_t, cursors_t, mask_t, batch_size, stack_size, batch_range):\n",
    "    \n",
    "    mask2 = T.cast(mask_t.dimshuffle(0, \"x\"), \"float32\")\n",
    "    top_next = mask2 * merge_value + (1 - mask2) * push_value\n",
    "    stack_next = T.set_subtensor(stack_t[t * batch_size + batch_range], top_next)\n",
    "    \n",
    "    # TODO grad_next\n",
    "    stack_grad_next = stack_grad_t\n",
    "    \n",
    "    cursors_next = cursors_t + (mask_t * -1 + (1 - mask_t) * 1)\n",
    "    queue_next = T.set_subtensor(queue_t[batch_range * stack_size + cursors_next], t)\n",
    "    \n",
    "    return stack_next, stack_grad_next, queue_next, cursors_next"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def step(t, x_t, mask_t,       # sequences\n",
    "         queue_t, cursors_t,   # outputs_info\n",
    "         stack_t, stack_grad_t # non_sequences\n",
    "        ):\n",
    "    \n",
    "    stack1 = stack_t[(t - 1) * batch_size + batch_range]\n",
    "    \n",
    "    stack2_ptrs = queue_t[cursors_t - 1 + batch_range * num_steps]\n",
    "    stack2 = stack_t[stack2_ptrs * batch_size + batch_range].reshape((batch_size, emb_dim))\n",
    "    \n",
    "    # Merge op: just add.\n",
    "    merge_value = stack1 + stack2\n",
    "    \n",
    "    stack_next, stack_grad_next, queue_next, cursors_next = stack_op(\n",
    "        t, stack_t, stack_grad_t, x_t, merge_value, queue_t, cursors_t,\n",
    "        mask_t, batch_size, num_steps, batch_range)\n",
    "    \n",
    "    outputs = (queue_next, cursors_next)\n",
    "    updates = {stack: stack_next,\n",
    "               stack_grad: stack_grad_next}\n",
    "    return outputs, updates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "scan_out, scan_updates = theano.scan(\n",
    "    step, [T.arange(num_steps), Xs, masks],\n",
    "    outputs_info=[queue, cursors],\n",
    "    non_sequences=[stack, stack_grad])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/jon/Projects/ipython-notebooks/.env/lib/python2.7/site-packages/ipykernel/__main__.py:2: UserWarning: The parameter 'updates' of theano.function() expects an OrderedDict, got <type 'dict'>. Using a standard dictionary here results in non-deterministic behavior. You should use an OrderedDict if you are using Python 2.7 (theano.compat.OrderedDict for older python), or use a list of (shared, update) pairs. Do not just convert your dictionary to this type before the call as the conversion will still be non-deterministic.\n",
      "  from ipykernel import kernelapp as app\n"
     ]
    }
   ],
   "source": [
    "f_clear = theano.function([], [], updates={stack: T.zeros_like(stack),\n",
    "                                           stack_grad: T.zeros_like(stack_grad)})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'theano.tensor.subtensor.IncSubtensor'>\n",
      "<class 'theano.tensor.subtensor.IncSubtensor'>\n",
      "<class 'theano.tensor.subtensor.IncSubtensor'>\n",
      "<class 'theano.tensor.subtensor.IncSubtensor'>\n",
      "IncSubtensor1 <class 'theano.tensor.subtensor.AdvancedIncSubtensor1'>\n",
      "\tAdvancedIncSubtensor1{no_inplace,set}(<TensorType(int32, vector)>, <TensorType(int64, scalar)>, Elemwise{add,no_inplace}.0)\n",
      "\tAdvancedIncSubtensor1{inplace,set}.0\n",
      "IncSubtensor1 <class 'theano.tensor.subtensor.AdvancedIncSubtensor1'>\n",
      "\tAdvancedIncSubtensor1{no_inplace,set}(stack_copy, Elemwise{Composite{((i0 * (i1 + i2)) + i3)}}.0, <TensorType(int64, vector)>)\n",
      "\tAdvancedIncSubtensor1{inplace,set}.0\n"
     ]
    }
   ],
   "source": [
    "f_ff = theano.function([Xs, masks], stack, updates=scan_updates)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "X_test = np.array([np.arange(batch_size)[:, np.newaxis].repeat(emb_dim, axis=1) + 1\n",
    "                   for t in range(num_steps)], dtype=np.int32)\n",
    "mask_test = np.array([[0, 0, 1, 0, 1],\n",
    "                      [0, 0, 0, 1, 0],\n",
    "                      [0, 0, 0, 0, 0]], dtype=np.int32).T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 1.,  1.,  1.,  1.,  1.,  1.,  1.],\n",
       "       [ 2.,  2.,  2.,  2.,  2.,  2.,  2.],\n",
       "       [ 3.,  3.,  3.,  3.,  3.,  3.,  3.],\n",
       "       [ 1.,  1.,  1.,  1.,  1.,  1.,  1.],\n",
       "       [ 2.,  2.,  2.,  2.,  2.,  2.,  2.],\n",
       "       [ 3.,  3.,  3.,  3.,  3.,  3.,  3.],\n",
       "       [ 2.,  2.,  2.,  2.,  2.,  2.,  2.],\n",
       "       [ 2.,  2.,  2.,  2.,  2.,  2.,  2.],\n",
       "       [ 3.,  3.,  3.,  3.,  3.,  3.,  3.],\n",
       "       [ 1.,  1.,  1.,  1.,  1.,  1.,  1.],\n",
       "       [ 4.,  4.,  4.,  4.,  4.,  4.,  4.],\n",
       "       [ 3.,  3.,  3.,  3.,  3.,  3.,  3.],\n",
       "       [ 3.,  3.,  3.,  3.,  3.,  3.,  3.],\n",
       "       [ 2.,  2.,  2.,  2.,  2.,  2.,  2.],\n",
       "       [ 3.,  3.,  3.,  3.,  3.,  3.,  3.]], dtype=float32)"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f_clear()\n",
    "f_ff(X_test, mask_test)\n",
    "stack.get_value()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}