tutorial_training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4d02d51dcadfcfb",
   "metadata": {},
   "source": [
    "# Tutorial: Training a model\n",
    "In this tutorial, we will train an event-based state-space model on a reduced version of the [Spiking Heidelberg Digits](https://zenkelab.org/resources/spiking-heidelberg-datasets-shd/) dataset.\n",
    "For training on larger datasets or multiple GPUs, we recommend using the training script `run_training.py` instead.\n",
    "\n",
    "## Setup\n",
    "\n",
    "Install and load the important modules and configuration. To install required packages, please do \n",
    "```\n",
    "pip3 install requirements.txt\n",
    "```\n",
    "\n",
    "Directories for loading datasets, model checkpoints and saving results are defined in the configuration file `system/local.yaml`.\n",
    "Please set your directories accordingly."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df7ee68ff3e429ed",
   "metadata": {},
   "source": [
    "## Data loading\n",
    "The SHD dataset contains 20 classes, digits from 0 to 9 in both German and English. \n",
    "We will use a reduced version of the dataset containing only two digits to train the model to non-trivial performance in reasonable time even on CPUs.\n",
    "\n",
    "[Download the training and test dataset](https://zenkelab.org/datasets/) and unpack the archives to `./data/`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f9883d23c86e5bcd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import Dataset, DataLoader, random_split\n",
    "import h5py\n",
    "import numpy as np\n",
    "\n",
    "class SpikingHeidelbergDigits(Dataset):\n",
    "    def __init__(self, path_to_file):\n",
    "        self.num_classes = 2\n",
    "        self.num_channels = 700\n",
    "        self.path_to_file = path_to_file\n",
    "        \n",
    "        # load the dataset\n",
    "        with h5py.File(path_to_file, 'r') as f:\n",
    "            self.channels = f['spikes']['units'][:]\n",
    "            self.timesteps = f['spikes']['times'][:]\n",
    "            self.labels = f['labels'][:]\n",
    "        \n",
    "        # filter the dataset to contain only two classes\n",
    "        mask = (self.labels == 0) | (self.labels == 1)\n",
    "        self.channels = self.channels[mask]\n",
    "        self.timesteps = self.timesteps[mask]\n",
    "        self.labels = self.labels[mask]\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.labels)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        # create tonic-like structured arrays\n",
    "        dtype = np.dtype([(\"t\", int), (\"x\", int), (\"p\", int)])\n",
    "        struct_arr = np.empty_like(self.channels[idx], dtype=dtype)\n",
    "        \n",
    "        # yield timesteps in milliseconds\n",
    "        timesteps = self.timesteps[idx] * 1e6\n",
    "        \n",
    "        struct_arr['t'] = timesteps\n",
    "        struct_arr['x'] = self.channels[idx]\n",
    "        struct_arr['p'] = 1\n",
    "        \n",
    "        # one-hot encoding of labels (required for CutMix augmentation)\n",
    "        label = np.eye(self.num_classes)[self.labels[idx]].astype(np.int32)\n",
    "            \n",
    "        return struct_arr, label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3be0429979f96a3f",
   "metadata": {},
   "outputs": [
    {
     "ename": "FileNotFoundError",
     "evalue": "[Errno 2] Unable to open file (unable to open file: name = 'data/shd_train.h5', errno = 2, error message = 'No such file or directory', flags = 0, o_flags = 0)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mFileNotFoundError\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[2], line 2\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;66;03m# Load the training and test dataset\u001b[39;00m\n\u001b[0;32m----> 2\u001b[0m train_dataset \u001b[38;5;241m=\u001b[39m \u001b[43mSpikingHeidelbergDigits\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mdata/shd_train.h5\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m      3\u001b[0m test_dataset \u001b[38;5;241m=\u001b[39m SpikingHeidelbergDigits(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mdata/shd_test.h5\u001b[39m\u001b[38;5;124m'\u001b[39m)\n",
      "Cell \u001b[0;32mIn[1], line 12\u001b[0m, in \u001b[0;36mSpikingHeidelbergDigits.__init__\u001b[0;34m(self, path_to_file)\u001b[0m\n\u001b[1;32m      9\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpath_to_file \u001b[38;5;241m=\u001b[39m path_to_file\n\u001b[1;32m     11\u001b[0m \u001b[38;5;66;03m# load the dataset\u001b[39;00m\n\u001b[0;32m---> 12\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m \u001b[43mh5py\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mFile\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpath_to_file\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mr\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m)\u001b[49m \u001b[38;5;28;01mas\u001b[39;00m f:\n\u001b[1;32m     13\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mchannels \u001b[38;5;241m=\u001b[39m f[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mspikes\u001b[39m\u001b[38;5;124m'\u001b[39m][\u001b[38;5;124m'\u001b[39m\u001b[38;5;124munits\u001b[39m\u001b[38;5;124m'\u001b[39m][:]\n\u001b[1;32m     14\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtimesteps \u001b[38;5;241m=\u001b[39m f[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mspikes\u001b[39m\u001b[38;5;124m'\u001b[39m][\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mtimes\u001b[39m\u001b[38;5;124m'\u001b[39m][:]\n",
      "File \u001b[0;32m/opt/homebrew/Caskroom/mambaforge/base/envs/ssm/lib/python3.12/site-packages/h5py/_hl/files.py:562\u001b[0m, in \u001b[0;36mFile.__init__\u001b[0;34m(self, name, mode, driver, libver, userblock_size, swmr, rdcc_nslots, rdcc_nbytes, rdcc_w0, track_order, fs_strategy, fs_persist, fs_threshold, fs_page_size, page_buf_size, min_meta_keep, min_raw_keep, locking, alignment_threshold, alignment_interval, meta_block_size, **kwds)\u001b[0m\n\u001b[1;32m    553\u001b[0m     fapl \u001b[38;5;241m=\u001b[39m make_fapl(driver, libver, rdcc_nslots, rdcc_nbytes, rdcc_w0,\n\u001b[1;32m    554\u001b[0m                      locking, page_buf_size, min_meta_keep, min_raw_keep,\n\u001b[1;32m    555\u001b[0m                      alignment_threshold\u001b[38;5;241m=\u001b[39malignment_threshold,\n\u001b[1;32m    556\u001b[0m                      alignment_interval\u001b[38;5;241m=\u001b[39malignment_interval,\n\u001b[1;32m    557\u001b[0m                      meta_block_size\u001b[38;5;241m=\u001b[39mmeta_block_size,\n\u001b[1;32m    558\u001b[0m                      \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwds)\n\u001b[1;32m    559\u001b[0m     fcpl \u001b[38;5;241m=\u001b[39m make_fcpl(track_order\u001b[38;5;241m=\u001b[39mtrack_order, fs_strategy\u001b[38;5;241m=\u001b[39mfs_strategy,\n\u001b[1;32m    560\u001b[0m                      fs_persist\u001b[38;5;241m=\u001b[39mfs_persist, fs_threshold\u001b[38;5;241m=\u001b[39mfs_threshold,\n\u001b[1;32m    561\u001b[0m                      fs_page_size\u001b[38;5;241m=\u001b[39mfs_page_size)\n\u001b[0;32m--> 562\u001b[0m     fid \u001b[38;5;241m=\u001b[39m \u001b[43mmake_fid\u001b[49m\u001b[43m(\u001b[49m\u001b[43mname\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmode\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43muserblock_size\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfapl\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfcpl\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mswmr\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mswmr\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    564\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(libver, \u001b[38;5;28mtuple\u001b[39m):\n\u001b[1;32m    565\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_libver \u001b[38;5;241m=\u001b[39m libver\n",
      "File \u001b[0;32m/opt/homebrew/Caskroom/mambaforge/base/envs/ssm/lib/python3.12/site-packages/h5py/_hl/files.py:235\u001b[0m, in \u001b[0;36mmake_fid\u001b[0;34m(name, mode, userblock_size, fapl, fcpl, swmr)\u001b[0m\n\u001b[1;32m    233\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m swmr \u001b[38;5;129;01mand\u001b[39;00m swmr_support:\n\u001b[1;32m    234\u001b[0m         flags \u001b[38;5;241m|\u001b[39m\u001b[38;5;241m=\u001b[39m h5f\u001b[38;5;241m.\u001b[39mACC_SWMR_READ\n\u001b[0;32m--> 235\u001b[0m     fid \u001b[38;5;241m=\u001b[39m \u001b[43mh5f\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mopen\u001b[49m\u001b[43m(\u001b[49m\u001b[43mname\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mflags\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfapl\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mfapl\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    236\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m mode \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mr+\u001b[39m\u001b[38;5;124m'\u001b[39m:\n\u001b[1;32m    237\u001b[0m     fid \u001b[38;5;241m=\u001b[39m h5f\u001b[38;5;241m.\u001b[39mopen(name, h5f\u001b[38;5;241m.\u001b[39mACC_RDWR, fapl\u001b[38;5;241m=\u001b[39mfapl)\n",
      "File \u001b[0;32mh5py/_objects.pyx:54\u001b[0m, in \u001b[0;36mh5py._objects.with_phil.wrapper\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32mh5py/_objects.pyx:55\u001b[0m, in \u001b[0;36mh5py._objects.with_phil.wrapper\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32mh5py/h5f.pyx:102\u001b[0m, in \u001b[0;36mh5py.h5f.open\u001b[0;34m()\u001b[0m\n",
      "\u001b[0;31mFileNotFoundError\u001b[0m: [Errno 2] Unable to open file (unable to open file: name = 'data/shd_train.h5', errno = 2, error message = 'No such file or directory', flags = 0, o_flags = 0)"
     ]
    }
   ],
   "source": [
    "# Load the training and test dataset\n",
    "train_dataset = SpikingHeidelbergDigits('data/shd_train.h5')\n",
    "test_dataset = SpikingHeidelbergDigits('data/shd_test.h5')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf72529578541b9",
   "metadata": {},
   "source": [
    "Check the length of the datasets to check if the data loading was successful."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec059aefa5d3408",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Number of training samples: {len(train_dataset)}\")\n",
    "print(f\"Number of test samples: {len(test_dataset)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee746256534411df",
   "metadata": {},
   "source": [
    "Now, create a validation set by randomly splitting the training dataset, and create data loaders for training, validation, and test datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ab24a1d63c4c194",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Split the training dataset into training and validation\n",
    "train_dataset, val_dataset = random_split(train_dataset, [int(0.8*len(train_dataset)), len(train_dataset) - int(0.8*len(train_dataset))])\n",
    "\n",
    "# Create data loaders\n",
    "from event_ssm.dataloading import event_stream_collate_fn\n",
    "from functools import partial\n",
    "\n",
    "collate_fn = partial(event_stream_collate_fn, resolution=(700,), pad_unit=8192)\n",
    "train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, drop_last=True, collate_fn=collate_fn)\n",
    "val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, drop_last=True, collate_fn=collate_fn)\n",
    "test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, collate_fn=collate_fn)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "acc88e3270fda10b",
   "metadata": {},
   "source": [
    "## Model definition\n",
    "We use the [hydra](https://hydra.cc/docs/intro/) package for efficient configuration management. Define the model configuration in a config file in the `configs` directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "810edc0798ad7622",
   "metadata": {},
   "outputs": [],
   "source": [
    "from hydra import compose, initialize\n",
    "from omegaconf import OmegaConf, open_dict\n",
    "\n",
    "with initialize(version_base=None, config_path=\"configs\", job_name=\"training tutorial\"):\n",
    "    cfg = compose(config_name=\"base\", overrides=[\"task=tutorial\"])\n",
    "\n",
    "with open_dict(cfg):    \n",
    "    # optax updates the schedule every iteration and not every epoch\n",
    "    cfg.optimizer.total_steps = cfg.training.num_epochs * len(train_loader) // cfg.optimizer.accumulation_steps\n",
    "    cfg.optimizer.warmup_steps = cfg.optimizer.warmup_epochs * len(train_loader) // cfg.optimizer.accumulation_steps\n",
    "    \n",
    "    # scale learning rate by batch size\n",
    "    cfg.optimizer.ssm_lr = cfg.optimizer.ssm_base_lr * cfg.training.per_device_batch_size * cfg.optimizer.accumulation_steps\n",
    "\n",
    "print(OmegaConf.to_yaml(cfg))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ca62d33ebabfdb2",
   "metadata": {},
   "source": [
    "Now, create the model using the configuration defined above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83e2062ea8b4fe02",
   "metadata": {},
   "outputs": [],
   "source": [
    "from event_ssm.ssm import init_S5SSM\n",
    "from event_ssm.seq_model import BatchClassificationModel\n",
    "\n",
    "ssm_init_fn = init_S5SSM(**cfg.model.ssm_init)\n",
    "model = BatchClassificationModel(\n",
    "    ssm=ssm_init_fn,\n",
    "    num_classes=test_dataset.num_classes,\n",
    "    num_embeddings=test_dataset.num_channels,\n",
    "    **cfg.model.ssm,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c855737447f70896",
   "metadata": {},
   "source": [
    "\n",
    "Initialize the training state by feeding a dummy input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4def1c65952a8ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "from event_ssm.train_utils import init_model_state\n",
    "\n",
    "# pick the first batch from the training loader\n",
    "batch = next(iter(train_loader))\n",
    "inputs, targets, timesteps, lengths = batch\n",
    "\n",
    "# initialize the training state\n",
    "key = jax.random.PRNGKey(cfg.seed)\n",
    "state = init_model_state(key, model, inputs, timesteps, lengths, cfg.optimizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "424bce6010abb8f1",
   "metadata": {},
   "source": [
    "## Inspect the model\n",
    "The model parameters are accessible as part of the training state. \n",
    "We will look into the spectrum of the recurrent operator here.\n",
    "The model was initialized with a single stage of blocks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a1602ed4a962265",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_spectrum(state):\n",
    "    params = state.params['encoder']['stages_0']\n",
    "    lambda_bar = []\n",
    "    time_scales = []\n",
    "    for name, sequence_layer in params.items():\n",
    "        # read lambda parameters\n",
    "        Lambda_im = sequence_layer['S5SSM_0']['Lambda_im']\n",
    "        Lambda_re = sequence_layer['S5SSM_0']['Lambda_re']\n",
    "        \n",
    "        # read and compute delta and Lambda\n",
    "        delta = np.exp(sequence_layer['S5SSM_0']['log_step'][:, 0])\n",
    "        Lambda = Lambda_re + 1j * Lambda_im\n",
    "        \n",
    "        # compute lambda_bar and time scales\n",
    "        lambda_bar.append(np.exp(Lambda * delta))\n",
    "        time_scales.append(1 / np.abs(Lambda) / delta)\n",
    "    return lambda_bar, time_scales\n",
    "spectrum, time_scales = get_spectrum(state)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d1377b48ebf3728",
   "metadata": {},
   "source": [
    "Plot the spectrum of the recurrent operator and the corresponding time scales upon initialization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ff987be826d7314",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "def plot_spectrum(spectrum):\n",
    "    fig, axes = plt.subplots(1, 6, figsize=(len(spectrum) * 4, 4))\n",
    "    # draw the unit circle\n",
    "    theta = np.linspace(0, 2 * np.pi, 100)  # 100 points from 0 to 2*pi\n",
    "    x = np.cos(theta)\n",
    "    y = np.sin(theta)\n",
    "    \n",
    "    # plot the spectrum\n",
    "    for i, (ax, layer) in enumerate(zip(axes, spectrum)):\n",
    "        ax.plot(x, y, 'r', linewidth=1)\n",
    "        ax.scatter(np.real(layer), np.imag(layer), marker='o', alpha=0.8)\n",
    "    \n",
    "        # format axis\n",
    "        ax.set_title(f'Layer {i}')\n",
    "        ax.set_aspect('equal', adjustable='box')\n",
    "        ax.set_xlim(-1.1, 1.1)\n",
    "        ax.set_ylim(-1.1, 1.1)\n",
    "    \n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "    \n",
    "plot_spectrum(spectrum)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1779ae6f5f72b44",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_time_scales(time_scales):\n",
    "    log_scales = np.log2(np.stack(time_scales).flatten())\n",
    "    fig, ax = plt.subplots(1, 1, figsize=(6, 6))\n",
    "    ax.hist(log_scales)\n",
    "    \n",
    "    # format axis\n",
    "    max_scale = np.max(np.ceil(log_scales))\n",
    "    min_scale = np.min(np.floor(log_scales))\n",
    "    ax.set_xlim((min_scale, max_scale))\n",
    "    xticks = np.arange(1 + max_scale - min_scale) + min_scale\n",
    "    ax.set_xticks(xticks, (2 ** xticks).astype(np.int32))\n",
    "    ax.set_title('Distribution of time scales')\n",
    "    ax.set_xlabel('Time scale')\n",
    "    ax.set_ylabel('Count')\n",
    "    plt.show()\n",
    "    \n",
    "plot_time_scales(time_scales)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4970c69f459df1d",
   "metadata": {},
   "source": [
    "## Train the model\n",
    "For training, we implemented a trainer module that makes training as easy as possible. The trainer module hides some boilerplate code for training from the user and provides a simple interface to train the model. It loops through the data loader, computes the loss, and updates the model parameters. Therefore, we need to define training_step and validation_step functions that the loop calls upon the model. These are implemented already, and can be used here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61ab72052c47f47c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from event_ssm.train_utils import training_step, evaluation_step\n",
    "from event_ssm.trainer import TrainerModule\n",
    "\n",
    "# just-in-time compile the training and evaluation functions\n",
    "train_step = jax.jit(training_step)\n",
    "eval_step = jax.jit(evaluation_step)\n",
    "\n",
    "# initialize the trainer module\n",
    "num_devices = 1\n",
    "trainer = TrainerModule(\n",
    "    train_state=state,\n",
    "    training_step_fn=train_step,\n",
    "    evaluation_step_fn=eval_step,\n",
    "    world_size=num_devices,\n",
    "    config=cfg,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d66a413bc0ac7d2b",
   "metadata": {},
   "source": [
    "We are now ready to start the training loop. \n",
    "\n",
    "**Note:** JAX compiles your program just-in-time (JIT) to optimize performance. This means that the first iteration of the training loop will be slower than the following ones.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d5ab8aa623db697",
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate random key for dropout\n",
    "key, dropout_key = jax.random.split(key)\n",
    "\n",
    "# train the model\n",
    "trainer.train_model(\n",
    "    train_loader=train_loader,\n",
    "    val_loader=val_loader,\n",
    "    test_loader=test_loader,\n",
    "    dropout_key=dropout_key\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a929b74f8ce235e5",
   "metadata": {},
   "source": [
    "## Inspect the trained model\n",
    "We now have a trained toy model on the SHD dataset.\n",
    "Let's look into the spectrum of the recurrent operator after training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3281a08743303429",
   "metadata": {},
   "outputs": [],
   "source": [
    "spectrum, time_scales = get_spectrum(trainer.train_state)\n",
    "plot_spectrum(spectrum)\n",
    "plot_time_scales(time_scales)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28ad17b8c7b61230",
   "metadata": {},
   "source": [
    "## Assignment\n",
    "The function `apply_ssm` in `event_ssm/ssm.py` implements the recurrent operator with an associative scan. On highly parallel GPUs, this can speed up training on very long sequences. \n",
    "On CPUs however, the overhead of the scan operation can slow down training. \n",
    "Your task is to implement a CPU-friendly version of the recurrent operator in `event_ssm/ssm.py` and compare the training time with the original implementation.\n",
    "We suggest to implement a step-by-step recurrence with [`jax.lax.scan`](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.scan.html) instead of the currenlty used [`jax.lax.associative_scan`](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.associative_scan.html) for this purpose."
   ]
  }
 ],
 "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.12.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}