annotatedS5.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# The Annotated S5"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Welcome! In this tutorial we will be implementing Simplified Structured State Models for Sequence Modeling, or more simply, S5, in Deepmind's Haiku neural network library for JAX. I've chosen to highlight the S5 because, though building off the original S4 and its variations, as the extra 'S' suggests, it provides a nice simplification of this family of architectures. If you are unfamiliar with any of these previous models I highly recommend checking out Sasha Rush's excellent blog post on the subject: https://srush.github.io/annotated-s4/, and its corresponding repository: https://github.com/srush/annotated-s4, in which he implements the original S4 in Flax. These were both extremely helpful for me in figuring out how to implement these architectures and serves as the direct inspiration for this notebook, so please do check them out! I have, however, tried to make this notebook as self-contained and accessible as possible so if you feel like jumping straight in, feel free. I should also mention the original repository from the paper authors: https://github.com/lindermanlab/S5, which itself builds upon Sasha Rush's Flax implementation of the S4s and has been the guiding light for my implementation. I have chosen to write mine in Haiku rather than Flax simply because that is my personal framework of choice, so if you are unfamiliar with Haiku, hopefully this can serve as a little introduction to that as well!"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "I will begin by briefly addressing what has been seen by many to be the black box of this family of architectures, and that is the High-order Polynomial Projection Operators, or HiPPOs, this is a new recurrent memory mechanism first introduced in https://arxiv.org/abs/2008.07669 that replaces the standard memory unit found in GRUs and LSTMs with an online function approximation based on a series of orthogonal polynomials, thereby creating an internal representation/memory of the function being modeled. A short comparison to standard RNN memory units is given in the paper mentioned above as follows:\n",
    "\n",
    "    \"by stacking multiple units in parallel and choosing a specific update function, we obtain the GRU update cell as a special case. In contrast to HiPPO which uses one hidden feature and projects it onto high order polynomials, these models use many hidden features but only project them with degree 1. This view sheds light on these classic techniques by showing how they can be derived from first principles.\"\n",
    "\n",
    "There's alot of fascinating components that go into these memory structures which I plan on writing another tutorial about, but, for now, I will leave it at the short description above and note that we will be initializing our HiPPOs with a little library called Hippox which I created for that very purpose, check out the source code if you feel like diving a bit deeper: https://github.com/JPGoodale/hippox (shameless self-promotion), and of course I highly recommend checking out the original papers and their code at https://github.com/HazyResearch/state-spaces."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Let's go ahead and import hippox and the core JAX libraries we will be using:\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import haiku as hk\n",
    "from hippox.main import Hippo\n",
    "from typing import Optional"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": true
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "First we'll define some helper functions for discretization and timescale initialization as the SSM equation is naturally continuous and must be made discrete to be unrolled as a linear recurrence like standard RNNs."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Here we are just using the zero order hold method for its sheer simplicity, with A, B and delta_t denoting the\n",
    "# state matrix, input matrix and change in timescale respectively.\n",
    "\n",
    "def discretize(A, B, delta_t):\n",
    "    Identity = jnp.ones(A.shape[0])\n",
    "    _A = jnp.exp(A*delta_t)\n",
    "    _B = (1/A * (_A-Identity)) * B\n",
    "    return _A, _B\n",
    "\n",
    "# This is a function used to initialize the trainable timescale parameter.\n",
    "def log_step_initializer(dt_min=0.001, dt_max=0.1):\n",
    "    def init(shape, dtype):\n",
    "        uniform = hk.initializers.RandomUniform()\n",
    "        return uniform(shape, dtype)*(jnp.log(dt_max) - jnp.log(dt_min)) + jnp.log(dt_min)\n",
    "    return init\n",
    "\n",
    "# Taken directly from https://github.com/deepmind/dm-haiku/blob/main/haiku/_src/recurrent.py\n",
    "def add_batch(nest, batch_size: Optional[int]):\n",
    "    broadcast = lambda x: jnp.broadcast_to(x, (batch_size,) + x.shape)\n",
    "    return jax.tree_util.tree_map(broadcast, nest)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": true
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "The linear SSM equation is as follows:\n",
    "        $ x0(t) = Ax(t) + Bu(t) $\n",
    "        $ y(t) = Cx(t) + Du(t) $\n",
    "\n",
    " We will now implement it as a recurrent Haiku module:"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "class LinearSSM(hk.RNNCore):\n",
    "    def __init__(self, state_size: int, name: Optional[str] = None):\n",
    "        super(LinearSSM, self).__init__(name=name)\n",
    "        # We won't get into the basis measure families here, just note that they are basically just the\n",
    "        # type of orthogonal polynomial we initialize with, the scaled Legendre measure (LegS) introduced\n",
    "        # in the original HiPPO paper is pretty much the standard initialization and is what is used in the\n",
    "        # main experiments in the S5 paper. I will also note that the Hippo class uses the diagonal representation\n",
    "        # of the state matrix by default, as this has become the standard in neural SSMs since shown to be\n",
    "        # equally effective as the diagonal plus low rank representation in https://arxiv.org/abs/2203.14343\n",
    "        # and then formalized in https://arxiv.org/abs/2206.11893.\n",
    "\n",
    "        _hippo = Hippo(state_size=state_size, basis_measure='legs')\n",
    "        # Must be called for parameters to be initialized\n",
    "        _hippo()\n",
    "\n",
    "        # We register the real and imaginary components of the state matrix A as separate parameters because\n",
    "        # they will have separate gradients in training, they will be conjoined back together and then discretized\n",
    "        # but this will simply be backpropagated through as a transformation of the lambda real and imaginary\n",
    "        # parameters (lambda is just what we call the diagonalized state matrix).\n",
    "\n",
    "        self._lambda_real = hk.get_parameter(\n",
    "            'lambda_real',\n",
    "            shape=[state_size,],\n",
    "            init=_hippo.lambda_initializer('real')\n",
    "        )\n",
    "        self._lambda_imag = hk.get_parameter(\n",
    "            'lambda_imaginary',\n",
    "            shape=[state_size,],\n",
    "            init=_hippo.lambda_initializer('imaginary')\n",
    "        )\n",
    "        self._A = self._lambda_real + 1j * self._lambda_imag\n",
    "\n",
    "       # For now, these initializations of the input and output matrices B and C match the S4D\n",
    "        # parameterization for demonstration purposes, we will implement the S5 versions later.\n",
    "\n",
    "        self._B = hk.get_parameter(\n",
    "            'B',\n",
    "            shape=[state_size,],\n",
    "            init=_hippo.b_initializer()\n",
    "        )\n",
    "        self._C = hk.get_parameter(\n",
    "            'C',\n",
    "            shape=[state_size, 2],\n",
    "            init=hk.initializers.RandomNormal(stddev=0.5**0.5)\n",
    "        )\n",
    "        self._output_matrix = self._C[..., 0] + 1j * self._C[..., 1]\n",
    "\n",
    "        # This feed-through matrix basically acts as a residual connection.\n",
    "        self._D = hk.get_parameter(\n",
    "            'D',\n",
    "            [1,],\n",
    "            init=jnp.ones,\n",
    "        )\n",
    "\n",
    "        self._delta_t = hk.get_parameter(\n",
    "            'delta_t',\n",
    "            shape=[1,],\n",
    "            init=log_step_initializer()\n",
    "        )\n",
    "        timescale = jnp.exp(self._delta_t)\n",
    "\n",
    "        self._state_matrix, self._input_matrix = discretize(self._A, self._B, timescale)\n",
    "\n",
    "    def __call__(self, inputs, prev_state):\n",
    "        u = inputs[:, jnp.newaxis]\n",
    "        new_state = self._state_matrix @ prev_state + self._input_matrix @ u\n",
    "        y_s = self._output_matrix @ new_state\n",
    "        out = y_s.reshape(-1).real + self._D * u\n",
    "        return out, new_state\n",
    "\n",
    "    def initial_state(self, batch_size: Optional[int] = None):\n",
    "        state = jnp.zeros([self._state_size])\n",
    "        if batch_size is not None:\n",
    "            state = add_batch(state, batch_size)\n",
    "        return state"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": true
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "You may notice that this looks an awful lot like a vanilla RNN cell, just with our special parameterization and without any activations, hence being a linear recurrence. I have initialized it as an instance of Haiku's RNN.Core abstract base class so that it can be unrolled using either the hk.dynamic_unroll or hk.static_unroll functions like any other recurrent module, however, if you are familiar with any of the S4 models you may be noticing that there's something crucial missing here: the convolutional representation. One of the key contributions of the S4 paper was its demonstration that the SSM ODE can be represented as either a linear recurrence, as above, for efficient inference, or as a global convolution for much faster training. That paper and the following papers then go on to present various kernels for efficiently computing this convolution with Fast Fourier Transforms, highly improving the computational efficiency of the model. Then why have we omitted them? Because the S5 architecture which we are about to explore simplifies all this by providing a purely recurrent representation in both training and inference, it does this by using a parallel recurrence that actually looks alot like a convolution itself! From the paper:\n",
    "\n",
    "    \"We use parallel scans to efficiently compute the states of a discretized linear SSM. Given a binary associative operator • (i.e. (a • b) • c = a • (b • c)) and a sequence of L elements [a1, a2, ..., aL], the scan operation (sometimes referred to as all-prefix-sum) returns the sequence [a1, (a1 • a2), ..., (a1 • a2 • ... • aL)].\"\n",
    "\n",
    "Let's see what this looks like in code, taken straight from the original author's implementation:"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "@jax.vmap\n",
    "def binary_operator(q_i, q_j):\n",
    "    A_i, b_i = q_i\n",
    "    A_j, b_j = q_j\n",
    "    return A_j * A_i, A_j * b_i + b_j\n",
    "\n",
    "def parallel_scan(A, B, C, inputs):\n",
    "    A_elements = A * jnp.ones((inputs.shape[0], A.shape[0]))\n",
    "    Bu_elements = jax.vmap(lambda u: B @ u)(inputs)\n",
    "    # Jax's built-in associative scan really comes in handy here as it executes a similar scan\n",
    "    # operation as used in a normal recurrent unroll but is specifically tailored to fit an associative\n",
    "    # operation like the one described in the paper.\n",
    "    _, xs = jax.lax.associative_scan(binary_operator, (A_elements, Bu_elements))\n",
    "    return jax.vmap(lambda x: (C @ x).real)(xs)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": true
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "It's that simple! In the original S4 we would have had to apply an independent singe-input, single-output (SISO) SSM for each feature of the input sequence such as in this excerpt from Sasha Rush's Flax implementation:"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def cloneLayer(layer):\n",
    "    return flax.linen.vmap(\n",
    "        layer,\n",
    "        in_axes=1,\n",
    "        out_axes=1,\n",
    "        variable_axes={\"params\": 1, \"cache\": 1, \"prime\": 1},\n",
    "        split_rngs={\"params\": True},\n",
    "    )\n",
    "SSMLayer = cloneLayer(SSMLayer)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": true
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Whereas in the S5 we process the entire sequence in one multi-input, multi-output (MIMO) layer."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Let's now rewrite our Module as a full S5 layer using this new method, we will be adding a few extra conditional arguments as well as changing some parameterization to match the original paper, but we'll walk through the reason for all these changes below."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# First we add a new helper function for the timescale initialization, this one just takes the previous\n",
    "# log_step_initializer and stores a bunch of them in an array since our model is now multi-in, multi-out.\n",
    "\n",
    "def init_log_steps(shape, dtype):\n",
    "    H = shape[0]\n",
    "    log_steps = []\n",
    "    for i in range(H):\n",
    "        log_step = log_step_initializer()(shape=(1,), dtype=dtype)\n",
    "        log_steps.append(log_step)\n",
    "\n",
    "    return jnp.array(log_steps)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": true
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# We will also rewrite our discretization for the MIMO context\n",
    "def discretize(A, B, delta_t):\n",
    "    Identity = jnp.ones(A.shape[0])\n",
    "    _A = jnp.exp(A*delta_t)\n",
    "    _B = (1/A * (_A-Identity))[..., None] * B\n",
    "    return _A, _B"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": true
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "class S5(hk.Module):\n",
    "    def __init__(self,\n",
    "                 state_size: int,\n",
    "\n",
    "                 # Now that we're MIMO we'll need to know the number of input features, commonly\n",
    "                 # referred to as the dimension of the model.\n",
    "                 d_model: int,\n",
    "\n",
    "                 # We must also now specify the number of blocks that we will split our matrices\n",
    "                 # into due to the MIMO context.\n",
    "                 n_blocks: int,\n",
    "\n",
    "                 # Short for conjugate symmetry, because our state matrix is complex we can half\n",
    "                 # the size of it since complex numbers are a real and imaginary number joined together,\n",
    "                 # this is not new to the S5, we just didn't mention it above.\n",
    "                 conj_sym: bool = True,\n",
    "\n",
    "                 # Another standard SSM argument that we omitted above for simplicity's sake,\n",
    "                 # this forces the real part of the state matrix to be negative for better\n",
    "                 # stability, especially in autoregressive tasks.\n",
    "                 clip_eigns: bool = False,\n",
    "\n",
    "                 # Like most RNNs, the S5 can be run in both directions if need be.\n",
    "                 bidirectional: bool = False,\n",
    "\n",
    "                 # Rescales delta_t for varying input resolutions, such as different audio\n",
    "                 # sampling rates.\n",
    "                 step_rescale: float = 1.0,\n",
    "                 name: Optional[str] = None\n",
    "    ):\n",
    "        super(S5, self).__init__(name=name)\n",
    "        self.conj_sym = conj_sym\n",
    "        self.bidirectional = bidirectional\n",
    "\n",
    "        # Note that the Hippo class takes conj_sym as an argument and will automatically half\n",
    "        # the state size provided in its initialization, which is why we need to provide a local\n",
    "        # state size that matches this for the shape argument in hk.get_parameter().\n",
    "\n",
    "        if conj_sym:\n",
    "            _state_size = state_size // 2\n",
    "        else:\n",
    "            _state_size = state_size\n",
    "\n",
    "        # With block_diagonal set as True and the number of blocks provided, our Hippo class\n",
    "        # will automatically handle this change of structure.\n",
    "\n",
    "        _hippo = Hippo(\n",
    "            state_size=state_size,\n",
    "            basis_measure='legs',\n",
    "            conj_sym=conj_sym,\n",
    "            block_diagonal=True,\n",
    "            n_blocks=n_blocks,\n",
    "        )\n",
    "        _hippo()\n",
    "\n",
    "        self._lambda_real = hk.get_parameter(\n",
    "            'lambda_real',\n",
    "            [_state_size],\n",
    "            init=_hippo.lambda_initializer('real')\n",
    "        )\n",
    "        self._lambda_imag = hk.get_parameter(\n",
    "            'lambda_imaginary',\n",
    "            [_state_size],\n",
    "            init=_hippo.lambda_initializer('imaginary')\n",
    "        )\n",
    "        if clip_eigns:\n",
    "            self._lambda = jnp.clip(self._lambda_real, None, -1e-4) + 1j * self._lambda_imag\n",
    "        else:\n",
    "            self._A = self._lambda_real + 1j * self._lambda_imag\n",
    "\n",
    "        # If you recall, I mentioned above that we are automatically using a diagonalized version of\n",
    "        # the HiPPO state matrix rather than the pure one, due to it being very hard to efficiently\n",
    "        # compute. I will now go into a little more detail on how this diagonal representation is\n",
    "        # derived, as it is important for how we initialize the input and output matrices. The diagonal\n",
    "        # decomposition of our state matrix is based on equivalence relation on the SSM parameters:\n",
    "        # (A, B, C) ∼ (V−1AV ,V−1B, CV) with V being the eigenvector of our original A matrix and V-1\n",
    "        # being the inverse eigenvector. The Hippo class has already performed the decomposition of A\n",
    "        # into (V-1AV) automatically, but we have not yet performed the decomposition of B and C, we will\n",
    "        # use the eigenvector_transform class method for that below, but first we must initialize B and C\n",
    "        # as normal distributions, lecun normal and truncated normal respectively. I will note that there\n",
    "        # are a few other options provided for C in the original repository but, to keep it simple, we will\n",
    "        # just use one here.\n",
    "\n",
    "        b_init = hk.initializers.VarianceScaling()\n",
    "        b_shape = [state_size, d_model]\n",
    "        b_init = b_init(b_shape, dtype=jnp.complex64)\n",
    "        self._B = hk.get_parameter(\n",
    "            'B',\n",
    "            [_state_size, d_model, 2],\n",
    "            init=_hippo.eigenvector_transform(b_init,  concatenate=True),\n",
    "        )\n",
    "        B = self._B[..., 0] + 1j * self._B[..., 1]\n",
    "\n",
    "        c_init = hk.initializers.TruncatedNormal()\n",
    "        c_shape = [d_model, state_size, 2]\n",
    "        c_init = c_init(c_shape, dtype=jnp.complex64)\n",
    "        self._C = hk.get_parameter(\n",
    "            'C',\n",
    "            [d_model, _state_size, 2],\n",
    "            init=_hippo.eigenvector_transform(c_init, inverse=False, concatenate=True),\n",
    "        )\n",
    "        # We need two output heads if bidirectional is True.\n",
    "        if bidirectional:\n",
    "            self._C2 = hk.get_parameter(\n",
    "                'C2',\n",
    "                [d_model, _state_size, 2],\n",
    "                init=_hippo.eigenvector_transform(c_init, inverse=False, concatenate=True),\n",
    "            )\n",
    "            C1 = self._C[..., 0] + 1j * self._C[..., 1]\n",
    "            C2 = self._C2[..., 0] + 1j * self._C2[..., 1]\n",
    "            self._output_matrix = jnp.concatenate((C1, C2), axis=-1)\n",
    "        else:\n",
    "            self._output_matrix = self._C[..., 0] + 1j * self._C[..., 1]\n",
    "\n",
    "        self._D = hk.get_parameter(\n",
    "            'D',\n",
    "            [d_model,],\n",
    "            init=hk.initializers.RandomNormal(stddev=1.0)\n",
    "        )\n",
    "\n",
    "        self._delta_t = hk.get_parameter(\n",
    "            'delta_T',\n",
    "            [_state_size, 1],\n",
    "            init=init_log_steps\n",
    "        )\n",
    "        timescale = step_rescale * jnp.exp(self._delta_t[:, 0])\n",
    "\n",
    "        # We could also use the bilinear discretization method, but we'll just stick to zoh for now.\n",
    "        self._state_matrix, self._input_matrix = discretize(self._A, B, timescale)\n",
    "\n",
    "\n",
    "    def __call__(self, inputs):\n",
    "        # Note that this is the exact same function as presented above just with alternate procedures\n",
    "        # depending on the bidirectional and conjugate symmetry arguments\n",
    "\n",
    "        A_elements = self._state_matrix * jnp.ones((inputs.shape[0], self._state_matrix.shape[0]))\n",
    "        Bu_elements = jax.vmap(lambda u: self._input_matrix @ u)(inputs)\n",
    "\n",
    "        _, xs = jax.lax.associative_scan(binary_operator, (A_elements, Bu_elements))\n",
    "\n",
    "        if self.bidirectional:\n",
    "            _, xs2 = jax.lax.associative_scan(binary_operator,\n",
    "                                          (A_elements, Bu_elements),\n",
    "                                          reverse=True)\n",
    "            xs = jnp.concatenate((xs, xs2), axis=-1)\n",
    "\n",
    "        if self.conj_sym:\n",
    "            ys = jax.vmap(lambda x: 2*(self._output_matrix @ x).real)(xs)\n",
    "        else:\n",
    "            ys = jax.vmap(lambda x: (self._output_matrix @ x).real)(xs)\n",
    "\n",
    "        Du = jax.vmap(lambda u: self._D * u)(inputs)\n",
    "\n",
    "        return ys + Du"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": true
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "There we have it, a complete S5 layer! Now let's form a block around it using a structure very similar to a transformer block with a Gated Linear Unit (GLU)."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "import dataclasses\n",
    "\n",
    "@dataclasses.dataclass\n",
    "class S5Block(hk.Module):\n",
    "    ssm: S5\n",
    "    d_model: int\n",
    "    dropout_rate: float\n",
    "    prenorm: bool\n",
    "    istraining: bool = True\n",
    "    name: Optional[str] = None\n",
    "\n",
    "    def __post_init__(self):\n",
    "        super(S5Block, self).__post_init__()\n",
    "        # We could use either layer norm or batch norm.\n",
    "        self._norm = hk.LayerNorm(axis=-1, create_scale=True, create_offset=True)\n",
    "        self._linear = hk.Linear(self.d_model)\n",
    "\n",
    "    def __call__(self, x):\n",
    "        skip = x\n",
    "        if self.prenorm:\n",
    "            x = self._norm(x)\n",
    "\n",
    "        x = self.ssm(x)\n",
    "        # There are a couple of other GLU patterns we could use here, but once again I have chosen\n",
    "        # one semi-arbitrarily to avoid cluttering our module with if statements.\n",
    "        x1 = hk.dropout(hk.next_rng_key(), self.dropout_rate, jax.nn.gelu(x))\n",
    "        x = x * jax.nn.sigmoid(self._linear(x1))\n",
    "        x = hk.dropout(hk.next_rng_key(), self.dropout_rate, x)\n",
    "\n",
    "        x = skip + x\n",
    "        if not self.prenorm:\n",
    "            x = self._norm(x)\n",
    "\n",
    "        return x"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": true
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Now let's make a stack of these blocks:"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "@dataclasses.dataclass\n",
    "class S5Stack(hk.Module):\n",
    "    ssm: S5\n",
    "    d_model: int\n",
    "    n_layers: int\n",
    "    dropout_rate: float\n",
    "    prenorm: bool\n",
    "    istraining: bool = True\n",
    "    name: Optional[str] = None\n",
    "\n",
    "    def __post_init__(self):\n",
    "        super(S5Stack, self).__post_init__(name=self.name)\n",
    "        self._encoder = hk.Linear(self.d_model)\n",
    "        self._layers = [\n",
    "            S5Block(\n",
    "                ssm=self.ssm,\n",
    "                d_model=self.d_model,\n",
    "                dropout_rate=self.dropout_rate,\n",
    "                istraining=self.istraining,\n",
    "                prenorm=self.prenorm,\n",
    "            )\n",
    "            for _ in range(self.n_layers)\n",
    "        ]\n",
    "\n",
    "    def __call__(self, x):\n",
    "        x = self._encoder(x)\n",
    "        for layer in self._layers:\n",
    "            x = layer(x)\n",
    "        return x"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": true
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "And finally a classifier on top, as we will be doing a simple classification task for this tutorial:"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "@dataclasses.dataclass\n",
    "class S5Classifier(hk.Module):\n",
    "    ssm: S5\n",
    "    d_model: int\n",
    "    d_output: int\n",
    "    n_layers: int\n",
    "    dropout_rate: float\n",
    "    mode: str = 'pool'\n",
    "    prenorm: bool = True\n",
    "    istraining: bool = True\n",
    "    name: Optional[str] = None\n",
    "\n",
    "    def __post_init__(self):\n",
    "        super(S5Classifier, self).__post_init__(name=self.name)\n",
    "        self._encoder = S5Stack(\n",
    "            ssm=self.ssm,\n",
    "            d_model=self.d_model,\n",
    "            n_layers=self.n_layers,\n",
    "            dropout_rate=self.dropout_rate,\n",
    "            istraining=self.istraining,\n",
    "            prenorm=self.prenorm,\n",
    "        )\n",
    "        self._decoder = hk.Linear(self.d_output)\n",
    "\n",
    "    def __call__(self, x):\n",
    "        x = self._encoder(x)\n",
    "        if self.mode == 'pool':\n",
    "            x = jnp.mean(x, axis=0)\n",
    "        elif self.mode == 'last':\n",
    "            x = x[-1]\n",
    "        else:\n",
    "            raise NotImplementedError(\"Mode must be in ['pool', 'last]\")\n",
    "        x = self._decoder(x)\n",
    "        return jax.nn.log_softmax(x, axis=-1)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": true
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Our model is now ready for training! Let's load some data. We will use the classic MNIST benchmark, but unlike the standard CNNs usually used for the task, we will be processing the images as a 1-dimensional sequence."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "from torch.utils.data import DataLoader\n",
    "from typing import NamedTuple\n",
    "\n",
    "# Code taken directly from https://github.com/srush/annotated-s4\n",
    "def create_mnist_classification_dataset(batch_size=128):\n",
    "    print(\"[*] Generating MNIST Classification Dataset...\")\n",
    "\n",
    "    # The usual 28*28 format of the images is now being flattened into a sequence of 784 pixels.\n",
    "    SEQ_LENGTH, N_CLASSES, IN_DIM = 784, 10, 1\n",
    "    tf = transforms.Compose(\n",
    "        [\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Normalize(mean=0.5, std=0.5),\n",
    "            transforms.Lambda(lambda x: x.view(IN_DIM, SEQ_LENGTH).t()),\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    train = torchvision.datasets.MNIST(\n",
    "        \"./data\", train=True, download=True, transform=tf\n",
    "    )\n",
    "    test = torchvision.datasets.MNIST(\n",
    "        \"./data\", train=False, download=True, transform=tf\n",
    "    )\n",
    "\n",
    "    trainloader = DataLoader(\n",
    "        train, batch_size, shuffle=True\n",
    "    )\n",
    "    testloader = DataLoader(\n",
    "        test, batch_size, shuffle=False\n",
    "    )\n",
    "\n",
    "    return trainloader, testloader, N_CLASSES, SEQ_LENGTH, IN_DIM\n",
    "\n",
    "Datasets = {\n",
    "    \"mnist-classification\": create_mnist_classification_dataset,\n",
    "}\n",
    "\n",
    "# Simple class for storing our dataset parameters.\n",
    "class Dataset(NamedTuple):\n",
    "    trainloader: DataLoader\n",
    "    testloader: DataLoader\n",
    "    n_classes: int\n",
    "    seq_length: int\n",
    "    d_input: int\n",
    "    classification: bool\n",
    "\n",
    "def create_dataset(dataset: str, batch_size: int) -> Dataset:\n",
    "    classification = 'classification' in dataset\n",
    "    dataset_init = Datasets[dataset]\n",
    "    trainloader, testloader, n_classes, seq_length, d_input = dataset_init(batch_size)\n",
    "    return Dataset(trainloader, testloader, n_classes, seq_length, d_input, classification)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": true
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Next we will set some hyperparameters:"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# I have halved the size of most of these parameters from what is used in the\n",
    "# paper for easier compute.\n",
    "STATE_SIZE: int = 128\n",
    "D_MODEL: int = 64\n",
    "N_LAYERS: int = 3\n",
    "N_BLOCKS: int = 4\n",
    "EPOCHS: int = 50\n",
    "BATCH_SIZE: int = 64\n",
    "DROPOUT_RATE: float = 0.1\n",
    "LEARNING_RATE: float = 1e-3\n",
    "SEED = 0"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": true
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "To keep things simple for this example, we will just be using a plain adam optimizer, but the results can be highly improved with extra techniques such as cosine annealing, learning rate schedules and weight decay."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "import optax\n",
    "\n",
    "class TrainingState(NamedTuple):\n",
    "    params: hk.Params\n",
    "    opt_state: optax.OptState\n",
    "    rng_key: jnp.ndarray\n",
    "\n",
    "optim = optax.adam(LEARNING_RATE)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": true
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "\n",
    "Now for the loss and update functions:"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "from functools import partial\n",
    "from typing import Tuple, MutableMapping, Any\n",
    "_Metrics = MutableMapping[str, Any]\n",
    "\n",
    "@partial(jnp.vectorize, signature=\"(c),()->()\")\n",
    "def cross_entropy_loss(logits, label) -> jnp.ndarray:\n",
    "    one_hot_label = jax.nn.one_hot(label, num_classes=logits.shape[0])\n",
    "    return -jnp.sum(one_hot_label * logits)\n",
    "\n",
    "\n",
    "@partial(jnp.vectorize, signature=\"(c),()->()\")\n",
    "def compute_accuracy(logits, label):\n",
    "    return jnp.argmax(logits) == label\n",
    "\n",
    "\n",
    "@partial(jax.jit, static_argnums=(3, 4))\n",
    "def update(\n",
    "        state: TrainingState,\n",
    "        inputs: jnp.ndarray,\n",
    "        targets: jnp.ndarray,\n",
    "        model: hk.transform,\n",
    "        classification: bool,\n",
    ") -> Tuple[TrainingState, _Metrics]:\n",
    "\n",
    "    rng_key, next_rng_key = jax.random.split(state.rng_key)\n",
    "\n",
    "    def loss_fn(params):\n",
    "        logits = model.apply(params, rng_key, inputs)\n",
    "        _loss = jnp.mean(cross_entropy_loss(logits, targets))\n",
    "        _accuracy = jnp.mean(compute_accuracy(logits, targets))\n",
    "        return _loss, _accuracy\n",
    "\n",
    "    if not classification:\n",
    "        targets = inputs[:, :, 0]\n",
    "\n",
    "    grad_fn = jax.value_and_grad(loss_fn, has_aux=True)\n",
    "    (loss, accuracy), gradients = grad_fn(state.params)\n",
    "    updates, new_opt_state = optim.update(gradients, state.opt_state, state.params)\n",
    "    new_params = optax.apply_updates(state.params, updates)\n",
    "\n",
    "    new_state = TrainingState(\n",
    "        params=new_params,\n",
    "        opt_state=new_opt_state,\n",
    "        rng_key=next_rng_key,\n",
    "    )\n",
    "    metrics = {\n",
    "        'loss': loss,\n",
    "        'accuracy': accuracy\n",
    "    }\n",
    "\n",
    "    return new_state, metrics\n",
    "\n",
    "\n",
    "@partial(jax.jit, static_argnums=(3, 4))\n",
    "def evaluate(\n",
    "        state: TrainingState,\n",
    "        inputs: jnp.ndarray,\n",
    "        targets: jnp.ndarray,\n",
    "        model: hk.transform,\n",
    "        classification,\n",
    ") -> _Metrics:\n",
    "\n",
    "    rng_key, _ = jax.random.split(state.rng_key, 2)\n",
    "\n",
    "    if not classification:\n",
    "        targets = inputs[:, :, 0]\n",
    "\n",
    "    logits = model.apply(state.params, rng_key, inputs)\n",
    "    loss = jnp.mean(cross_entropy_loss(logits, targets))\n",
    "    accuracy = jnp.mean(compute_accuracy(logits, targets))\n",
    "\n",
    "    metrics = {\n",
    "        'loss': loss,\n",
    "        'accuracy': accuracy\n",
    "    }\n",
    "\n",
    "    return metrics"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": true
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "\n",
    "Now we call these update and evaluate functions in their respective epochs:"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "def training_epoch(\n",
    "        state: TrainingState,\n",
    "        trainloader: DataLoader,\n",
    "        model: hk.transform,\n",
    "        classification: bool = False,\n",
    ") -> Tuple[TrainingState, jnp.ndarray, jnp.ndarray]:\n",
    "\n",
    "    batch_losses, batch_accuracies = [], []\n",
    "    for batch_idx, (inputs, targets) in enumerate(tqdm(trainloader)):\n",
    "        inputs = jnp.array(inputs.numpy())\n",
    "        targets = jnp.array(targets.numpy())\n",
    "        state, metrics = update(\n",
    "            state, inputs, targets,\n",
    "            model, classification\n",
    "        )\n",
    "        batch_losses.append(metrics['loss'])\n",
    "        batch_accuracies.append(metrics['accuracy'])\n",
    "\n",
    "    return (\n",
    "        state,\n",
    "        jnp.mean(jnp.array(batch_losses)),\n",
    "        jnp.mean(jnp.array(batch_accuracies))\n",
    "    )\n",
    "\n",
    "\n",
    "def validation_epoch(\n",
    "        state: TrainingState,\n",
    "        testloader: DataLoader,\n",
    "        model: hk.transform,\n",
    "        classification: bool = True,\n",
    ") -> Tuple[jnp.ndarray, jnp.ndarray]:\n",
    "\n",
    "    losses, accuracies = [], []\n",
    "    for batch_idx, (inputs, targets) in enumerate(tqdm(testloader)):\n",
    "        inputs = jnp.array(inputs.numpy())\n",
    "        targets = jnp.array(targets.numpy())\n",
    "        metrics = evaluate(\n",
    "            state, inputs, targets,\n",
    "            model, classification\n",
    "        )\n",
    "        losses.append(metrics['loss'])\n",
    "        accuracies.append(metrics['accuracy'])\n",
    "\n",
    "    return jnp.mean(jnp.array(losses)), jnp.mean(jnp.array(accuracies))"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": true
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "import torch\n",
    "# Set random number generators\n",
    "torch.random.manual_seed(SEED)\n",
    "key = jax.random.PRNGKey(SEED)\n",
    "rng, init_rng = jax.random.split(key)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": true
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Create our dataset and dummy data for initialization of the model's params.\n",
    "ds = create_dataset('mnist-classification', BATCH_SIZE)\n",
    "init_data = jnp.array(next(iter(ds.trainloader))[0].numpy())"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": true
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# In Haiku, we have to call our model inside a transformed function using hk.transform for it to become\n",
    "# functionally pure and compatible with essential JAX functions like jax.grad(). Here we are using hk.vmap\n",
    "# instead of jax.vmap because we are calling it from within a hk.transform.\n",
    "@hk.transform\n",
    "def forward(x) -> hk.transform:\n",
    "    neural_net = S5Classifier(\n",
    "        S5(\n",
    "            STATE_SIZE,\n",
    "            D_MODEL,\n",
    "            N_BLOCKS,\n",
    "        ),\n",
    "        D_MODEL,\n",
    "        ds.n_classes,\n",
    "        N_LAYERS,\n",
    "        DROPOUT_RATE,\n",
    "    )\n",
    "    return hk.vmap(neural_net, split_rng=False)(x)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": true
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Set state\n",
    "initial_params = forward.init(init_rng, init_data)\n",
    "initial_opt_state = optim.init(initial_params)\n",
    "\n",
    "state = TrainingState(\n",
    "    params=initial_params,\n",
    "    opt_state=initial_opt_state,\n",
    "    rng_key=rng\n",
    ")"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": true
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "And finally our training loop!"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "for epoch in range(EPOCHS):\n",
    "    print(f\"[*] Training Epoch {epoch + 1}...\")\n",
    "    state, training_loss, training_accuracy = training_epoch(\n",
    "        state,\n",
    "        ds.trainloader,\n",
    "        forward,\n",
    "        ds.classification\n",
    "    )\n",
    "    print(f\"[*] Running Epoch {epoch + 1} Validation...\")\n",
    "    test_loss, test_accuracy = validation_epoch(\n",
    "        state,\n",
    "        ds.testloader,\n",
    "        forward,\n",
    "        ds.classification\n",
    "    )\n",
    "    print(f\"\\n=>> Epoch {epoch + 1} Metrics ===\")\n",
    "    print(\n",
    "        f\"\\tTrain Loss: {training_loss:.5f} -- Train Accuracy:\"\n",
    "        f\" {training_accuracy:.4f}\\n\\t Test Loss: {test_loss:.5f} --  Test\"\n",
    "        f\" Accuracy: {test_accuracy:.4f}\"\n",
    "    )"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}