Experiment2.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nN2OaMWcczfa"
      },
      "source": [
        "# **Transformers learn in-context by gradient descent**\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FR8YNR-g9JXA"
      },
      "outputs": [],
      "source": [
        "#@title Imports external sources\n",
        "import os\n",
        "import io\n",
        "import PIL.Image, PIL.ImageDraw, PIL.ImageFont\n",
        "import base64\n",
        "import zipfile\n",
        "import json\n",
        "import requests\n",
        "import matplotlib.pylab as pl\n",
        "import numpy as np\n",
        "import glob\n",
        "import requests\n",
        "import random as pyrandom\n",
        "from concurrent import futures\n",
        "from functools import partial\n",
        "from scipy.ndimage import rotate\n",
        "from IPython.display import Image, HTML, clear_output\n",
        "from tqdm import tqdm_notebook, tnrange\n",
        "import time\n",
        "from typing import Any, MutableMapping, NamedTuple, Tuple\n",
        "\n",
        "!pip install --quiet --upgrade jax\n",
        "!pip install --quiet --upgrade jaxlib\n",
        "import jax\n",
        "from jax import grad, jit, vmap\n",
        "import jax.numpy as jnp\n",
        "\n",
        "!pip install --quiet -U dm-haiku\n",
        "!pip install --quiet -U optax\n",
        "import haiku as hk\n",
        "import math\n",
        "!pip install --quiet -U ml_collections\n",
        "from ml_collections import config_dict\n",
        "import matplotlib.pylab as pl\n",
        "import matplotlib.colors as mcolors\n",
        "colors = pl.colormaps['Dark2']"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "933ztM3DSREA"
      },
      "outputs": [],
      "source": [
        "!git clone --quiet https://github.com/aliayub40995/transformers-jax\n",
        "%cd /content/transformers-jax\n",
        "from src.transformer import Transformer\n",
        "from src.data import create_reg_data, create_weights\n",
        "from src.config import config\n",
        "from src.train import *"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4KwAI4LZFfcF"
      },
      "outputs": [],
      "source": [
        "#@title Config\n",
        "softmax = False\n",
        "layernorm = False\n",
        "\n",
        "recurrent_Transformer = True\n",
        "num_layers = 1\n",
        "num_heads = 1\n",
        "num_seeds = 1\n",
        "\n",
        "config.seed = 0\n",
        "config.local_usage = True\n",
        "\n",
        "####\n",
        "config.deq = recurrent_Transformer\n",
        "config.gd_deq = recurrent_Transformer\n",
        "config.att_only_trans = True\n",
        "####\n",
        "\n",
        "config.pre_train_gd = True\n",
        "config.train_gd_whitening = True\n",
        "config.train_gd_lr = True\n",
        "config.use_bias = False\n",
        "config.include_query = False\n",
        "\n",
        "config.distract_size = 0\n",
        "config.training_steps = 5000 if config.deq else 50000\n",
        "config.training_steps_gd = 1000 if config.gd_deq else 30000\n",
        "config.use_softmax = softmax\n",
        "config.first_layer_sm = False\n",
        "config.use_non_lin_mix = False\n",
        "\n",
        "config.widening_factor = 4\n",
        "\n",
        "config.layer_norm = layernorm\n",
        "config.out_proj = False\n",
        "config.in_proj = False\n",
        "config.adam = True\n",
        "config.dataset_size = 10\n",
        "config.input_size = 10\n",
        "config.key_size = 11\n",
        "config.num_layers = num_layers\n",
        "config.num_heads = num_heads\n",
        "\n",
        "config.grad_clip_value = 10 if num_layers > 2 else 0.001\n",
        "config.grad_clip_value_gd = 10 if num_layers > 2 else 0.001\n",
        "config.lr = 0.0005 if num_layers > 3 else 0.001\n",
        "config.gd_lr = 0.0005 if num_layers > 3 else 0.001\n",
        "config.wd = 0.0\n",
        "config.init_scale = 0.002 / config.num_layers\n",
        "config.bs = 2048\n",
        "config.bs_gd_train = 512\n",
        "config.dampening = 1.0\n",
        "config.clip = 10 if num_layers > 3 else 0\n",
        "\n",
        "config.dropout_rate = 0.0\n",
        "data_creator = vmap(create_reg_data,\n",
        "                    in_axes=(0, None, None, None, None, None),\n",
        "                    out_axes=0)\n",
        "\n",
        "config.y_update = False\n",
        "config.input_range = 1\n",
        "\n",
        "config.pos_enc = False\n",
        "config.pos_enc_size = 20\n",
        "config.concat_pos_enc = False\n",
        "config.analyse = True\n",
        "\n",
        "config.cycle_data = 0 #0 means online learning\n",
        "config.num_seeds = num_seeds\n",
        "if config.num_layers == 1:\n",
        "  assert config.deq == True\n",
        "  assert config.gd_deq == True\n",
        "\n",
        "if config.num_layers > 1:\n",
        "  assert config.y_update == False\n",
        "\n",
        "config.in_proj = False\n",
        "config.emb_size = 11"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "z9P7XZf7rh6l"
      },
      "outputs": [],
      "source": [
        "#@title Utils\n",
        "pl.rcParams.update({'font.size': 12})\n",
        "pl.rc('axes', labelsize=14)\n",
        "pl.rcParams.update({\n",
        "    \"text.usetex\": False,\n",
        "})\n",
        "\n",
        "import matplotlib.colors as mcolors\n",
        "colors = pl.colormaps['Dark2']\n",
        "def np2pil(a):\n",
        "  if a.dtype in [np.float32, np.float64]:\n",
        "    a = np.uint8(np.clip(a, 0, 1)*255)\n",
        "  return PIL.Image.fromarray(a)\n",
        "\n",
        "def imwrite(f, a, fmt=None):\n",
        "  a = np.asarray(a)\n",
        "  if isinstance(f, str):\n",
        "    fmt = f.rsplit('.', 1)[-1].lower()\n",
        "    if fmt == 'jpg':\n",
        "      fmt = 'jpeg'\n",
        "    f = open(f, 'wb') #GFile.open(f, 'wb')\n",
        "  np2pil(a).save(f, fmt, quality=95)\n",
        "\n",
        "def imencode(a, fmt='jpeg'):\n",
        "  a = np.asarray(a)\n",
        "  if len(a.shape) == 3 and a.shape[-1] == 4:\n",
        "    fmt = 'png'\n",
        "  f = io.BytesIO()\n",
        "  imwrite(f, a, fmt)\n",
        "  return f.getvalue()\n",
        "\n",
        "def moving_average(x, w):\n",
        "    return np.convolve(x, np.ones(w), 'valid') / w\n",
        "\n",
        "def grab_plot(close=True):\n",
        "  \"\"\"Return the current Matplotlib figure as an image.\"\"\"\n",
        "  fig = pl.gcf()\n",
        "  fig.canvas.draw()\n",
        "  img = np.array(fig.canvas.renderer._renderer)\n",
        "  a = np.float32(img[..., 3:]/255.0)\n",
        "  img = np.uint8(255*(1.0-a) + img[...,:3] * a)  # alpha\n",
        "  if close:\n",
        "    pl.close()\n",
        "  return img\n",
        "\n",
        "def display_learning(train, test=None, gt=None, inter=None, title=\"train\",\n",
        "                     title1=\"Trained TF\", title2=\"Test\",\n",
        "                     title3='Gradient descent', title4='Interpolated',\n",
        "                     y_label1 = 'L2 Norm', y_label2 = 'Cosine sim',\n",
        "                     y_lim_l=0,  y_lim_u=1, single_seeds= False,\n",
        "                     plot_title = None,\n",
        "                     y_lim_u2= 1., y_lim_l2=0.,  x_label = 'Training steps',\n",
        "                     second_axis=False, color_add=0, rw=10, num_iter_os=None,\n",
        "                     allow_download=False, plot_num=1, two_plots=False,\n",
        "                     loc_first = 'upper left', label_title=\"Loss\",\n",
        "                     loc_sec='upper left', yscale_log=False, line=\"-\",\n",
        "                     color_axis=True,\n",
        "                     height=3.5, width = 4, ax1=None, ax2=None):\n",
        "\n",
        "  \"\"\"Update learning curve image.\"\"\"\n",
        "\n",
        "  train_list = train\n",
        "  train = np.array(train)\n",
        "  num_seeds_train = train.shape[0]\n",
        "  train_std = np.std(train, axis=0)\n",
        "  train = np.mean(train, axis=0)\n",
        "\n",
        "  if test is not None:\n",
        "    test_list = test\n",
        "    test_std = np.std(test, axis=0)\n",
        "    test = np.mean(test, axis=0)\n",
        "\n",
        "  if gt is not None:\n",
        "    gt_list = gt\n",
        "    gt_std = np.std(gt, axis=0)\n",
        "    gt = np.mean(gt, axis=0)\n",
        "\n",
        "  if inter is not None:\n",
        "    inter_list = inter\n",
        "    inter_std = np.std(inter, axis=0)\n",
        "    inter = np.mean(inter, axis=0)\n",
        "\n",
        "  if plot_num == 1:\n",
        "    fig, ax1 = pl.subplots()\n",
        "    ax1.set_xlabel(x_label)\n",
        "    fig.set_size_inches(width, height)\n",
        "\n",
        "\n",
        "  if test is not None and not second_axis:\n",
        "    x_range = np.arange(0, num_iter_os, int(num_iter_os/len(test)))\n",
        "    if len(test_list) > 1:\n",
        "      if single_seeds:\n",
        "        for s in test_list:\n",
        "          ax1.plot(x_range, s, color=colors(0.1+color_add), alpha=0.2, label=title2,linewidth='2')\n",
        "      else:\n",
        "        ax1.fill_between(x_range, test-test_std, test+test_std ,alpha=0.2, facecolor=colors(0.1+color_add))\n",
        "    ax1.plot(x_range, test, color=colors(0.1+color_add), label=title2,linewidth='3')\n",
        "    #test_avg = moving_average(test, rw)\n",
        "    #ax1.plot(x_range[:len(test_avg)], test_avg, color=colors(0.1+color_add), label=title2)\n",
        "\n",
        "  if gt is not None:\n",
        "    if not second_axis:\n",
        "      x_range = np.arange(0, num_iter_os, int(num_iter_os/len(gt)))\n",
        "      #ax1.plot(x_range[:len(gt[:-rw])], gt[:-rw], color=colors(0.2+color_add), alpha=0.3)\n",
        "      #gt_avg = moving_average(gt, rw)\n",
        "      ax1.plot(x_range, gt, color=colors(0.2+color_add), label=title3,linewidth='3')\n",
        "      if len(gt_list) > 1:\n",
        "        if single_seeds:\n",
        "          for s in gt_list:\n",
        "            ax1.plot(x_range, s, color=colors(0.2+color_add), alpha=0.2, linewidth='2', zorder=0)\n",
        "        else:\n",
        "          ax1.fill_between(x_range, gt-gt_std, gt+gt_std,alpha=0.2, facecolor=colors(0.2+color_add))\n",
        "    else:\n",
        "      x_range = np.arange(0, num_iter_os, int(num_iter_os/len(gt)))\n",
        "      ax1.plot(x_range, gt, color=colors(0.6+color_add), label=title3,linewidth='3')\n",
        "      if len(gt_list) > 1:\n",
        "        if single_seeds:\n",
        "          for s in gt_list:\n",
        "            ax1.plot(x_range, s, color=colors(0.6+color_add), alpha=0.3, linewidth='2', zorder=0)\n",
        "        else:\n",
        "          ax1.fill_between(x_range, gt-gt_std, gt+gt_std ,alpha=0.2, facecolor=colors(0.6+color_add))\n",
        "\n",
        "  if test is not None and second_axis:\n",
        "    x_range = np.arange(0, num_iter_os, int(num_iter_os/len(test)))\n",
        "    ax1.plot(x_range, test, color=colors(0.5+color_add), label=title2,linewidth='3')\n",
        "    #test_avg = moving_average(test, rw)\n",
        "    #ax1.plot(x_range[:len(test_avg)],test_avg, color=colors(0.5+color_add))\n",
        "    if len(test_list) > 1:\n",
        "      if single_seeds:\n",
        "        for s in test_list:\n",
        "          ax1.plot(x_range, s, color=colors(0.5+color_add), linewidth='2', alpha=0.3, zorder=0)\n",
        "      else:\n",
        "        ax1.fill_between(x_range, test-test_std, test+test_std ,alpha=0.2, facecolor=colors(0.5+color_add))\n",
        "\n",
        "  if inter is not None and not second_axis:\n",
        "    x_range = np.arange(0, num_iter_os, int(num_iter_os/len(inter)))\n",
        "    ax1.plot(x_range, inter, color=colors(0.4+color_add), label=title4, linewidth='3', zorder=10)\n",
        "    if len(inter_list) > 1:\n",
        "      if single_seeds:\n",
        "        for s in inter_list:\n",
        "          ax1.plot(x_range, s, color=colors(0.4+color_add), alpha=0.3, linewidth='2', zorder=0)\n",
        "      else:\n",
        "        ax1.fill_between(x_range, inter-inter_std, inter+inter_std ,alpha=0.2, facecolor=colors(0.4+color_add), zorder=1)\n",
        "    #inter_avg = moving_average(inter, rw)\n",
        "    #ax1.plot(x_range[:len(inter_avg)], inter_avg, color=colors(0.7+color_add), label=title4)\n",
        "\n",
        "\n",
        "  if second_axis:\n",
        "    if ax2 is None:\n",
        "      ax2 = ax1.twinx()\n",
        "    ax2.set_zorder(0)\n",
        "    ax1.set_zorder(1)\n",
        "    ax1.set_frame_on(False)\n",
        "    #train_avg = moving_average(train, rw)\n",
        "    #ax2.plot(train[:-rw], color=colors(0.1+color_add), alpha=0.3)\n",
        "    ax2.plot(x_range, train, color=colors(0.4+color_add), label=title1, linewidth='3')\n",
        "    ax2.plot(x_range, np.ones_like(train), \"--\", color=\"gray\", linewidth='0.7')\n",
        "    if len(train_list) > 1:\n",
        "      if single_seeds:\n",
        "        for s in train_list:\n",
        "          print(x_range, s)\n",
        "          ax1.plot(x_range, s, line, color=colors(0.4+color_add), alpha=0.3, linewidth='2', zorder=0)\n",
        "      else:\n",
        "        ax2.fill_between(x_range, train-train_std, train+train_std ,alpha=0.2, facecolor=colors(0.4+color_add))\n",
        "\n",
        "    if color_axis:\n",
        "      ax2.yaxis.label.set_color(colors(0.4+color_add))\n",
        "    else:\n",
        "      legend2 = ax2.legend(loc='upper right', framealpha=0.99, facecolor='white')\n",
        "      legend2.set_zorder(100)\n",
        "    ax2.spines['top'].set_visible(False)\n",
        "  else:\n",
        "    #train_avg = moving_average(train, rw)\n",
        "    if line != \"-\":\n",
        "      ax1.scatter(x_range, train, s=[100 for _ in x_range],\n",
        "                  marker=\"+\", color=colors(0.3+color_add), alpha=1, label=title1, zorder=3, linewidths=3)\n",
        "    else:\n",
        "      ax1.plot(x_range, train, line, color=colors(0.3+color_add), label=title1, linewidth='3', zorder=11)\n",
        "    #ax1.plot(x_range[:len(train_avg)], train_avg, line, color=colors(0.3+color_add), label=title1)\n",
        "    if len(train_list) > 1:\n",
        "      if single_seeds:\n",
        "          for s in train_list:\n",
        "            ax1.plot(x_range, s, line, color=colors(0.3+color_add), alpha=0.3, linewidth='2', zorder=0)\n",
        "      else:\n",
        "        ax1.fill_between(x_range, train-train_std, train+train_std,\n",
        "                       alpha=0.5, facecolor=colors(0.3+color_add))\n",
        "\n",
        "    ax1.legend(loc='best', framealpha=1, facecolor='white')\n",
        "    ax1.spines['right'].set_visible(False)\n",
        "    legend = ax1.legend(loc='upper right', framealpha=0.99, facecolor='white')\n",
        "    legend.set_zorder(100)\n",
        "\n",
        "  legend1 = ax1.legend(loc=loc_first, framealpha=0.99, facecolor='white')\n",
        "  legend1.set_zorder(100)\n",
        "  if second_axis:\n",
        "    ax2.set_ylabel(y_label2)\n",
        "    ax1.set_ylabel(y_label1)\n",
        "    ax1.set_ylim(y_lim_l, y_lim_u)\n",
        "    legend1 = ax1.legend(loc=loc_sec, framealpha=0.99, facecolor='white')\n",
        "    ax2.set_ylim(y_lim_l2, y_lim_u2)\n",
        "    ax1.set_ylim(bottom=0)\n",
        "  else:\n",
        "    pl.ylabel(label_title)\n",
        "    pl.ylim(y_lim_l, y_lim_u)\n",
        "  ax1.spines['top'].set_visible(False)\n",
        "\n",
        "  if plot_title is not None:\n",
        "    pl.title(plot_title)\n",
        "\n",
        "  if yscale_log:\n",
        "    ax1.set_yscale(\"log\")\n",
        "  #pl.title(title)\n",
        "  pl.tight_layout()\n",
        "\n",
        "  if allow_download:\n",
        "    if second_axis:\n",
        "      pl.savefig(\"sim.pdf\", format=\"pdf\")\n",
        "      %download_file sim.pdf\n",
        "    else:\n",
        "      pl.savefig(\"train.pdf\", format=\"pdf\")\n",
        "      %download_file train.pdf\n",
        "  else:\n",
        "    img = grab_plot()\n",
        "    display(Image(data=imencode(img, fmt='jpeg')), display_id=title)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8ttht5tTEbUs"
      },
      "outputs": [],
      "source": [
        "#@title Lists\n",
        "\n",
        "loss_trans_list =  [[]  for _ in range(config.num_seeds)]\n",
        "loss_trans_train_list =  [[]  for _ in range(config.num_seeds)]\n",
        "losses_gd_list =  [[]  for _ in range(config.num_seeds)]\n",
        "losses_gd_list_trained =  [[]  for _ in range(config.num_seeds)]\n",
        "losses_int_list_trained =  [[]  for _ in range(config.num_seeds)]\n",
        "cos_sim_list, cos_sim_list_o =  [[]  for _ in range(config.num_seeds)],  [[]  for _ in range(config.num_seeds)]\n",
        "grad_norm_list, grad_norm_list_o =  [[]  for _ in range(config.num_seeds)],  [[]  for _ in range(config.num_seeds)]\n",
        "p_norm_list, p_norm_list_o =  [[]  for _ in range(config.num_seeds)],  [[]  for _ in range(config.num_seeds)]\n",
        "\n",
        "cos_sim_list, cos_sim_list_o =  [[]  for _ in range(config.num_seeds)],  [[]  for _ in range(config.num_seeds)]\n",
        "grad_norm_list, grad_norm_list_o =  [[]  for _ in range(config.num_seeds)],  [[]  for _ in range(config.num_seeds)]\n",
        "p_norm_list, p_norm_list_o =  [[]  for _ in range(config.num_seeds)],  [[]  for _ in range(config.num_seeds)]\n",
        "\n",
        "ir_t_list = [[]  for _ in range(config.num_seeds)]\n",
        "ws_t_list = [[]  for _ in range(config.num_seeds)]\n",
        "ir_gd_list = [[]  for _ in range(config.num_seeds)]\n",
        "ws_gd_list = [[]  for _ in range(config.num_seeds)]\n",
        "\n",
        "ir_t_ood_list = [[]  for _ in range(config.num_seeds)]\n",
        "ws_t_ood_list = [[]  for _ in range(config.num_seeds)]\n",
        "ir_gd_ood_list = [[]  for _ in range(config.num_seeds)]\n",
        "ws_gd_ood_list = [[]  for _ in range(config.num_seeds)]\n",
        "\n",
        "ir_gd_trained_list = [[]  for _ in range(config.num_seeds)]\n",
        "ws_gd_trained_list = [[]  for _ in range(config.num_seeds)]\n",
        "ir_gd_ood_trained_list = [[]  for _ in range(config.num_seeds)]\n",
        "ws_gd_ood_trained_list = [[]  for _ in range(config.num_seeds)]\n",
        "\n",
        "ir_inter_list = [[]  for _ in range(config.num_seeds)]\n",
        "ws_inter_list = [[]  for _ in range(config.num_seeds)]\n",
        "ir_inter_ood_list = [[]  for _ in range(config.num_seeds)]\n",
        "ws_inter_ood_list = [[]  for _ in range(config.num_seeds)]\n",
        "\n",
        "\n",
        "losses_noisy_list = [[]  for _ in range(config.num_seeds)]\n",
        "losses_gd_noisy_list = [[]  for _ in range(config.num_seeds)]\n",
        "losses_gd_noisy_trained_list = [[]  for _ in range(config.num_seeds)]\n",
        "losses_inter_noisy_list = [[]  for _ in range(config.num_seeds)]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "acU234z4zqLf"
      },
      "outputs": [],
      "source": [
        "#@title Logic how to interpolate weights\n",
        "def interpolate_weights(train_state, params_gd, deq=False):\n",
        "  if (config.num_heads == 1 and\n",
        "      config.sum_norm == False and config.deq == True and\n",
        "      config.layer_norm == False and config.att_only_trans == True):\n",
        "\n",
        "    cur_train_params = {k.replace('transformer', 'Transformer_gd'):v.copy() for\n",
        "                    k,v in train_state.params.items()}\n",
        "\n",
        "    inter_params = {k.replace('transformer', 'Transformer_gd'): {'w': jnp.zeros_like(v['w'])} for\n",
        "                    k,v in train_state.params.items()}\n",
        "\n",
        "    for k,v in cur_train_params.items():\n",
        "      if \"key\" in k:\n",
        "        key_gd  = params_gd[k]['w'].copy()\n",
        "        key  = cur_train_params[k]['w'].copy()\n",
        "      if \"linear\" in k:\n",
        "        linear_gd = params_gd[k]['w'].copy()\n",
        "        linear = cur_train_params[k]['w'].copy()\n",
        "      if \"query\" in k:\n",
        "        query_gd = params_gd[k]['w'].copy()\n",
        "        query = cur_train_params[k]['w'].copy()\n",
        "      if \"value\" in k:\n",
        "        value_gd = params_gd[k]['w'].copy()\n",
        "        value = cur_train_params[k]['w'].copy()\n",
        "\n",
        "        query = jnp.matmul(query, key.T)\n",
        "        #print(query)\n",
        "        key = jnp.identity(query.shape[0])\n",
        "        mean = np.mean([query[a, a] for a in range(query.shape[0]-1)])\n",
        "        query = query/mean\n",
        "        query_gd = jnp.matmul(query_gd, key.T)\n",
        "        key_gd = jnp.identity(query.shape[0])\n",
        "        query = (query + query_gd)/2\n",
        "\n",
        "        linear = jnp.matmul(value, linear)\n",
        "        #print(linear)\n",
        "        value = jnp.identity(query.shape[0])\n",
        "        linear = linear*mean\n",
        "        linear_gd = jnp.matmul(value_gd, linear_gd)\n",
        "        value_gd = jnp.identity(query.shape[0])\n",
        "        linear = (linear + linear_gd)/2\n",
        "\n",
        "        inter_params[k.replace('value', 'linear')]['w'] = linear\n",
        "        inter_params[k.replace('value', 'value')]['w'] = value\n",
        "        inter_params[k.replace('value', 'query')]['w'] = query\n",
        "        inter_params[k.replace('value', 'key')]['w'] = key\n",
        "\n",
        "    losses_int, _, _ = predict_test.apply(inter_params, eval_rng, eval_data, True)\n",
        "  else:\n",
        "    losses_int = None\n",
        "    inter_params = None\n",
        "  return losses_int, inter_params"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7s4jr3OBpaq9"
      },
      "outputs": [],
      "source": [
        "#@title Training\n",
        "\n",
        "# interpolate GD and trained TF\n",
        "inter = True if (config.deq and not config.use_softmax and config.num_heads == 1) else False\n",
        "\n",
        "eval_rng = jax.random.PRNGKey(5)\n",
        "for cur_seed in range(0, config.num_seeds):\n",
        "  config.seed = cur_seed\n",
        "  optimiser, train_state, _, rng = init()\n",
        "  rng, data_rng = jax.random.split(rng, 2)\n",
        "  if config.analyse:\n",
        "    lr_min, min_loss = scan_lrs(eval_rng, lin_diag=False, bs=10000)\n",
        "    if cur_seed == 0:\n",
        "      print('Best lr found for ', config.num_layers ,' steps of gradient descent: ', lr_min/config.dataset_size, \" with loss \", min_loss)\n",
        "\n",
        "    params_gd = create_weights(config.input_size, 1, config.dataset_size, lr_min,\n",
        "                                jax.random.normal(data_rng, shape=[1, 1, config.input_size])*0 ,\n",
        "                              lin_diag=False, gd_deq=config.gd_deq,\n",
        "                               num_layers=config.num_layers,\n",
        "                               input_mlp_rnd=rng if (config.input_mlp or config.in_proj) else None,\n",
        "                               in_proj=config.in_proj)\n",
        "    if config.num_layers > 1 or (config.in_proj and config.num_layers == 1):\n",
        "      if cur_seed == 0:\n",
        "        lr_min, min_loss = scan_lrs(eval_rng, lin_diag=True, bs=10000)\n",
        "        params_init = create_weights(config.input_size, 1, config.dataset_size, lr_min,\n",
        "                                  jax.random.normal(data_rng, shape=[1, 1, config.input_size])*0,\n",
        "                                lin_diag=True, gd_deq=config.gd_deq,\n",
        "                                num_layers=config.num_layers,\n",
        "                                input_mlp_rnd=eval_rng if (config.input_mlp or config.in_proj) else None,\n",
        "                               in_proj=config.in_proj)\n",
        "        params_gd_trained, data_rng = pre_train_gd_hps(eval_rng, params_init)\n",
        "    else:\n",
        "        params_gd_trained = params_gd\n",
        "\n",
        "  eval_data = data_creator(jax.random.split(eval_rng, num=10000),\n",
        "                               config.input_size,\n",
        "                               config.dataset_size,\n",
        "                               config.size_distract,\n",
        "                               config.input_range,\n",
        "                               config.weight_scale)\n",
        "  if config.analyse:\n",
        "    loss_gd, _, _ = predict_test.apply(params_gd, eval_rng, eval_data,True)\n",
        "    loss_gd_trained, _, _ = predict_test.apply(params_gd_trained, eval_rng,\n",
        "                                                eval_data, True)\n",
        "  original_data_rng = data_rng\n",
        "  for step in range(config.training_steps):\n",
        "    if config.cycle_data > 0:\n",
        "      if step % config.cycle_data == 0:\n",
        "        data_rng = original_data_rng\n",
        "\n",
        "    rng, data_rng = jax.random.split(data_rng, 2)\n",
        "    train_data = data_creator(jax.random.split(rng, num=config.bs),\n",
        "                              config.input_size,\n",
        "                              config.dataset_size,\n",
        "                              config.size_distract,\n",
        "                              config.input_range,\n",
        "                              config.weight_scale)\n",
        "    train_state, metrics = update(train_state, train_data, optimiser)\n",
        "    if step % 100 == 0:\n",
        "\n",
        "      loss_trans, _, _ = predict_test.apply(train_state.params, eval_rng,\n",
        "                                            eval_data, False)\n",
        "      loss_trans_list[cur_seed].append(loss_trans)\n",
        "      loss_trans_train_list[cur_seed].append(metrics['train_loss'].item(),)\n",
        "      if config.analyse:\n",
        "        losses_gd_list[cur_seed].append(loss_gd)\n",
        "        losses_gd_list_trained[cur_seed].append(loss_gd_trained)\n",
        "\n",
        "        losses_int, inter_params = interpolate_weights(train_state, params_gd_trained)\n",
        "\n",
        "        losses_int_list_trained[cur_seed].append(losses_int)\n",
        "\n",
        "        #rng, data_rng, eval_rng = jax.random.split(data_rng, 3)\n",
        "        # Alignment Transformers and GD\n",
        "        cos_sim, w_norm, p_norm = analyse(eval_data, train_state, eval_rng,\n",
        "                                          params_gd)\n",
        "        cos_sim_o, w_norm_o, p_norm_o = analyse(eval_data, train_state, eval_rng,\n",
        "                                          params_gd_trained)\n",
        "        if step > 0:\n",
        "          display((\"Current seed\", cur_seed,\n",
        "                   \"Training step\", step, \"Gradient descent loss\", loss_gd.item(),\n",
        "                    \"GD ++ loss\", loss_gd_trained.item(),\n",
        "                    \"Trained TF loss\", loss_trans.item(),\n",
        "                    \"Interpolated model loss\", losses_int.item() if inter else \"-\",\n",
        "                    \"Cosine sim TF vs GD\", cos_sim.item(),\n",
        "                    \"Cosine sim TF vs GD++\", cos_sim_o.item() if config.num_layers > 1 else \"-\"),\n",
        "                    display_id=\"Cur met\")\n",
        "\n",
        "        cos_sim_list[cur_seed].append(cos_sim)\n",
        "        grad_norm_list[cur_seed].append(w_norm)\n",
        "        p_norm_list[cur_seed].append(p_norm)\n",
        "\n",
        "        cos_sim_list_o[cur_seed].append(cos_sim_o)\n",
        "        grad_norm_list_o[cur_seed].append(w_norm_o)\n",
        "        p_norm_list_o[cur_seed].append(p_norm_o)\n",
        "\n",
        "      else:\n",
        "        print(step, loss_trans)\n",
        "\n",
        "  if config.analyse:\n",
        "\n",
        "    # Out-of-training-distribution behaviour\n",
        "    ir_t, ws_t, ir_gd, ws_gd, _ = ood(train_state, eval_rng, params_gd, 1000)\n",
        "\n",
        "    ir_t_list[cur_seed].append(ir_t)\n",
        "    ws_t_list[cur_seed].append(ws_t)\n",
        "    ir_gd_list[cur_seed].append(ir_gd)\n",
        "    ws_gd_list[cur_seed].append(ws_gd)\n",
        "\n",
        "    # More Out-of-training-distribution behaviour\n",
        "    ir_t_ood, ws_t_ood, ir_gd_ood, ws_gd_ood, _=ood_other_d(train_state,\n",
        "                                                                  eval_rng,\n",
        "                                                                  params_gd)\n",
        "    ir_t_ood_list[cur_seed].append(ir_t_ood)\n",
        "    ws_t_ood_list[cur_seed].append(ws_t_ood)\n",
        "    ir_gd_ood_list[cur_seed].append(ir_gd_ood)\n",
        "    ws_gd_ood_list[cur_seed].append(ws_gd_ood)\n",
        "\n",
        "    # Out-of-training-distribution behaviour\n",
        "    _, _, ir_gd_trained, ws_gd_trained, _ = ood(train_state, eval_rng,\n",
        "                                            params_gd_trained, 1000)\n",
        "\n",
        "    # More Out-of-training-distribution behaviour\n",
        "    _, _, ir_gd_ood_trained, ws_gd_ood_trained, _= ood_other_d(train_state,\n",
        "                                                                  eval_rng,\n",
        "                                                                  params_gd_trained)\n",
        "\n",
        "    ir_gd_trained_list[cur_seed].append(ir_gd_trained)\n",
        "    ws_gd_trained_list[cur_seed].append(ws_gd_trained)\n",
        "    ir_gd_ood_trained_list[cur_seed].append(ir_gd_ood_trained)\n",
        "    ws_gd_ood_trained_list[cur_seed].append(ws_gd_ood_trained)\n",
        "\n",
        "    if inter:\n",
        "      #Out-of-training-distribution behaviour\n",
        "      _, _, ir_inter, ws_inter, _ = ood(train_state, eval_rng,\n",
        "                                              inter_params, 1000)\n",
        "\n",
        "      #More Out-of-training-distribution behaviour\n",
        "      _, _, ir_inter_ood, ws_inter_ood, _= ood_other_d(train_state, eval_rng,\n",
        "                                                      inter_params)\n",
        "\n",
        "      ir_inter_list[cur_seed].append(ir_inter)\n",
        "      ws_inter_list[cur_seed].append(ws_inter)\n",
        "      ir_inter_ood_list[cur_seed].append(ir_inter_ood)\n",
        "      ws_inter_ood_list[cur_seed].append(ws_inter_ood)\n",
        "\n",
        "    # Noisy training-distribution behaviour\n",
        "    losses_noisy, losses_gd_noisy = noisy_data_ana(train_state, eval_rng,\n",
        "                                                  params_gd)\n",
        "    _, losses_gd_noisy_trained = noisy_data_ana(train_state, eval_rng,\n",
        "                                                  params_gd_trained)\n",
        "\n",
        "    losses_noisy_list[cur_seed].append(losses_noisy)\n",
        "    losses_gd_noisy_list[cur_seed].append(losses_gd_noisy)\n",
        "    losses_gd_noisy_trained_list[cur_seed].append(losses_gd_noisy_trained)\n",
        "\n",
        "    if inter:\n",
        "      _, losses_inter = noisy_data_ana(train_state, eval_rng, inter_params)\n",
        "      losses_inter_noisy_list[cur_seed].append(losses_inter)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tWjKF2SgrbTO"
      },
      "outputs": [],
      "source": [
        "#@title Visualize loss and alignment measures (the printing layout is set to default and might be suboptimal)\n",
        "\n",
        "cosine_low = 0.0\n",
        "if config.num_layers == 1:\n",
        "  display_learning(loss_trans_list, test=[losses_gd_list[0]], y_lim_u=0.4, y_lim_l=0.2,\n",
        "                  rw=1, title=\"train.pdf\", allow_download=False,\n",
        "                  single_seeds = True, label_title =\"Loss\",\n",
        "                  title2='GD', title1='Trained TF',\n",
        "                  title3='GD',  loc_first='upper right',\n",
        "                  num_iter_os=len(loss_trans_list[0])*100)\n",
        "\n",
        "  display_learning(cos_sim_list, grad_norm_list, p_norm_list,\n",
        "                  title1=\"Model cos\",\n",
        "                  title2=\"Model diff\", y_lim_u=2,\n",
        "                  title3=\"Preds diff\", second_axis=True, color_add=0.2,\n",
        "                  y_lim_u2=1.19, loc_sec='center right', single_seeds = False,\n",
        "                  y_lim_l2=cosine_low, color_axis=False, width= 5, y_label2 = 'Cosine sim',\n",
        "                  rw=1, num_iter_os=len(grad_norm_list[0])*100, title=\"sim.pdf\",\n",
        "                  allow_download=False)\n",
        "else:\n",
        "  display_learning(loss_trans_list, gt=losses_gd_list_trained,\n",
        "                   test=[losses_gd_list[0]], y_lim_u=0.3, y_lim_l=0.0,\n",
        "                   rw=1, title=\"train.pdf\", allow_download=False,\n",
        "                   title2='GD', title1='Trained TF',\n",
        "                   title3='GD$^{++}$',loc_first='upper right', x_label = \"Training steps\",\n",
        "                   single_seeds = True, plot_title = None,\n",
        "                   num_iter_os=len(loss_trans_list[0])*100)\n",
        "\n",
        "  display_learning(cos_sim_list, grad_norm_list, p_norm_list, title1=\"Model cos\",\n",
        "                  title2=\"Model diff\", y_lim_u=1.8,\n",
        "                  title3=\"Preds diff\", second_axis=True, color_add=0.2,\n",
        "                  y_lim_u2=1.09999999, color_axis=False,  width= 4, x_label = \"Training steps\",\n",
        "                  plot_title=\"GD vs trained TF\",\n",
        "                  y_lim_l2=0.5, loc_sec = 'center right', y_label1 = 'L2 Norm', y_label2 = 'Cosine sim',\n",
        "                  rw=1, num_iter_os=len(loss_trans_list[0])*100, title=\"sim.pdf\",\n",
        "                  allow_download=False, plot_num=1)\n",
        "\n",
        "  display_learning(cos_sim_list_o, grad_norm_list_o, p_norm_list_o, title1=\"Model cos\",\n",
        "                  title2=\"Model diff\", y_lim_u=1.8, x_label = \"Training steps\", plot_title = \"GD$^{++}$ vs trained TF\" ,\n",
        "                  title3=\"Preds diff\", second_axis=True, color_add=0.2,\n",
        "                  y_lim_u2=1.0599999, color_axis=False, width= 4, y_label1 = 'L2 Norm', y_label2 = 'Cosine sim',\n",
        "                  y_lim_l2=0.5, loc_sec = 'center right',\n",
        "                  rw=1, num_iter_os=len(loss_trans_list[0])*100, title=\"sim2.pdf\",\n",
        "                  allow_download=False)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Pxg3ohU9a7Re"
      },
      "outputs": [],
      "source": [
        "#@title Visualize iteratively applying a single layer Transformer layer\n",
        "if config.gd_deq and config.num_heads == 1 and config.num_layers == 1:\n",
        "  scale = 0.75\n",
        "\n",
        "  train_state_final = train_state\n",
        "  losses_gd, losses_tr, x_range = [], [], []\n",
        "\n",
        "  orig_damp = config.dampening\n",
        "  config.dampening = scale\n",
        "  orig_num_layers = config.num_layers\n",
        "  config.num_layers = 50\n",
        "\n",
        "  config.y_update = True\n",
        "  params_gd = create_weights(config.input_size, 1, config.dataset_size, lr_min*scale,\n",
        "                                jax.random.normal(data_rng, shape=[1, 1, config.input_size])*0,\n",
        "                              lin_diag=False,\n",
        "                              gd_deq=config.gd_deq, num_layers=config.num_layers)\n",
        "  x_range = np.arange(1, config.num_layers)\n",
        "\n",
        "  _, int_params = interpolate_weights(train_state, params_gd)\n",
        "\n",
        "  _, _, _, rng = init()\n",
        "  _, _, losses_gd = predict_test.apply(params_gd, eval_rng, eval_data, True)\n",
        "  _, _, losses_gd_train = predict_test.apply(params_gd_trained, eval_rng, eval_data, True)\n",
        "  _, _, losses_tr = predict_test.apply(train_state.params, eval_rng, eval_data, False)\n",
        "  _, _, losses_int = predict_test.apply(int_params, eval_rng, eval_data, True)\n",
        "  config.num_layers = orig_num_layers\n",
        "  config.orig_damp = orig_damp\n",
        "  config.y_update = False\n",
        "  train_state = train_state_final\n",
        "\n",
        "if config.gd_deq and config.deq and config.num_heads == 1:\n",
        "  if orig_num_layers > 1:\n",
        "      display_learning([losses_tr], gt=[losses_gd_train], test=[losses_gd],\n",
        "                      y_lim_u=10, y_lim_l=0.0,\n",
        "                    rw=1, title=\"extra_gd_steps.pdf\", allow_download=False,\n",
        "                    title3='GD$^{++}$',\n",
        "                    title2='GD', title1='Trained TF',\n",
        "                    line=\"*\", height = 2.5, loc_first='lower right',\n",
        "                    x_label = 'Gradient steps / TF  steps',\n",
        "                    num_iter_os=len(losses_tr))\n",
        "  else:\n",
        "      each = 3\n",
        "      display_learning([losses_tr[0::each]], test=[losses_gd[0::each]], y_lim_u=0.27, y_lim_l=0.00,\n",
        "                  rw=1, title=\"extra_gd_steps.pdf\", allow_download=False,\n",
        "                  #plot_title=\"Dampening $\\lambda =$\" +str(scale),\n",
        "                  title3='GD$^{++}$' if config.pre_train_gd else 'GD',\n",
        "                  title2='Gradient descent', title1='Trained Transformer',\n",
        "                  line=\"+\", height = 2.5, loc_first='upper right',\n",
        "                  x_label = 'GD Steps / Transformer Layers',\n",
        "                  num_iter_os=len(losses_tr[0::each])*each)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "C_XuqAMIVgCJ"
      },
      "outputs": [],
      "source": [
        "#@title Visualize Transformer performance on larger inputs (comment out for other OOD)\n",
        "\n",
        "#@title Config\n",
        "Test_on_larger_inputs = True\n",
        "Test_on_larger_targets = False\n",
        "same_distribution = True\n",
        "\n",
        "assert Test_on_larger_inputs or Test_on_larger_targets\n",
        "\n",
        "log = True\n",
        "every = 5\n",
        "l_bound, u_bound = 0.06, 10\n",
        "leg_loc = 'upper left'\n",
        "\n",
        "if Test_on_larger_inputs:\n",
        "\n",
        "  if same_distribution:\n",
        "    title = r'Test on larger inputs'\n",
        "\n",
        "    x_axis = r'$\\alpha$   where x $\\sim$ U($-\\alpha, \\alpha$)'\n",
        "    download = 'Input scale'\n",
        "    data = [ir_gd_list, ir_gd_trained_list, ir_t_list, ir_inter_list, np.arange(0.5, 2+0.03, 0.03)]\n",
        "    training_x = np.ones([100])\n",
        "    step = np.max(np.array(data[:1]))/100.\n",
        "    training_y = np.arange(0, np.max(np.array(data[:1])) - step, step)\n",
        "  else:\n",
        "\n",
        "    leg_loc = 'upper left'\n",
        "    x_axis = r'$\\alpha$   where $\\alpha x$'\n",
        "    l_bound, u_bound = 0.008, 10\n",
        "    every = 5\n",
        "    data = [ir_gd_ood_list, ir_gd_ood_trained_list, ir_t_ood_list, ir_inter_ood_list, np.arange(0.5, 5+0.05, 0.05)]\n",
        "    training_x = np.ones([100])\n",
        "    step = np.max(np.array(data[:1]))/100.\n",
        "    training_y = np.arange(0, np.max(np.array(data[:1])) - step, step)\n",
        "\n",
        "if Test_on_larger_targets:\n",
        "  title = r'Test on larger targets'\n",
        "\n",
        "  if same_distribution:\n",
        "    x_axis = r'$\\alpha W$   where W $\\sim$ N(0, I)'\n",
        "    download = 'Weight scale'\n",
        "    leg_loc = 'lower right'\n",
        "    data = [ws_gd_list, ws_gd_trained_list, ws_t_list, ws_inter_list, np.arange(0.5, 5+0.1, 0.1)]\n",
        "    training_x = np.ones([100])\n",
        "    step = np.max(np.array(data[:1]))/100.\n",
        "    training_y = np.arange(0, np.max(np.array(data[:1])) - step, step)\n",
        "\n",
        "  else:\n",
        "\n",
        "    x_axis = r'$\\alpha W$   where W $\\sim$ N(0, I)'\n",
        "    download = 'Weight scale ood'\n",
        "    data = [ws_gd_ood_list, ws_gd_ood_trained_list, ws_t_ood_list, ws_inter_ood_list, np.arange(0.5, 5+0.1, 0.1)]\n",
        "    training_x = np.ones([100])\n",
        "    step = np.max(np.array(data[:1]))/100.\n",
        "    training_y = np.arange(0, np.max(np.array(data[:1])) - step, step)\n",
        "\n",
        "pl.rcParams.update({'font.size': 12})\n",
        "\n",
        "fig, ax1 = pl.subplots()\n",
        "fig.set_size_inches(4, 3.5)\n",
        "pl.xlabel(x_axis)\n",
        "pl.ylabel('Loss')\n",
        "if log:\n",
        "  pl.yscale('log')\n",
        "\n",
        "#print(training_x, training_y)\n",
        "pl.plot(training_x[:len(training_y)], training_y, \"--\", color=\"gray\", linewidth='0.7',zorder=0)\n",
        "\n",
        "stretch = data[-1][0::every]\n",
        "gd_list = np.array(data[0])\n",
        "num_seeds = gd_list.shape[0]\n",
        "\n",
        "ir_gd_std = np.std(gd_list, axis=0)[0][0::every]\n",
        "ir_gd = np.mean(gd_list, axis=0)[0][0::every]\n",
        "\n",
        "ir_gd_trained_std = np.std(data[1], axis=0)[0][0::every]\n",
        "ir_gd_trained = np.mean(data[1], axis=0)[0][0::every]\n",
        "\n",
        "ir_t_std = np.std(data[2], axis=0)[0][0::every]\n",
        "ir_t = np.mean(data[2], axis=0)[0][0::every]\n",
        "\n",
        "ax1.scatter(stretch, ir_gd, s=[150 for _ in stretch], marker=\"v\", color=colors(0.1),\n",
        "         label=\"GD\",linewidth=3, zorder=2)\n",
        "ax1.plot(stretch, ir_gd, color=colors(0.1)\n",
        "         ,linewidth='2', zorder=2)\n",
        "\n",
        "#if num_seeds > 1:\n",
        "#  ax1.fill_between(stretch, ir_gd-ir_gd_std*1, ir_gd+ir_gd_std*1, alpha=0.2, facecolor=colors(0.1), zorder=1)\n",
        "\n",
        "if config.num_layers > 1:\n",
        "  ax1.scatter(stretch, ir_gd_trained, s=[150 for _ in stretch], marker=\"x\", color=colors(0.2),\n",
        "          label=\"GD$^{++}$\", zorder=4, linewidths=3)\n",
        "\n",
        "  ax1.plot(stretch, ir_gd_trained, color=colors(0.2), linewidth=2, zorder=4)\n",
        "  #if num_seeds > 1:\n",
        "  #  ax1.fill_between(stretch, ir_gd_trained-ir_gd_trained_std*1, ir_gd_trained+ir_gd_trained_std*1, alpha=0.2, facecolor=colors(0.2))\n",
        "\n",
        "if inter:\n",
        "  ir_inter_std = np.std(data[3], axis=0)[0][0::every]\n",
        "  ir_inter = np.mean(data[3], axis=0)[0][0::every]\n",
        "  ax1.scatter(stretch, ir_inter, s=[10 for _ in stretch], marker=\"o\", color=colors(0.4), label=\"Interpolated\", linewidths=5, zorder=7)\n",
        "  ax1.plot(stretch, ir_inter, color=colors(0.4), linewidth=2, zorder=7)\n",
        "  #ax1.plot(stretch, ir_inter, \"-\", color=colors(0.4))\n",
        "  #if num_seeds > 1:\n",
        "  #  ax1.fill_between(stretch, ir_inter-ir_inter_std*1, ir_inter+ir_inter_std*1, alpha=0.2, facecolor=colors(0.4), zorder=1)\n",
        "\n",
        "ax1.scatter(stretch, ir_t, s=[230 for _ in stretch], marker=\"+\", color=colors(0.3), alpha=1, label=\"Trained TF\", zorder=5, linewidths=3)\n",
        "#ax1.plot(stretch, ir_t, \"-\", color=colors(0.3))\n",
        "#if num_seeds > 1:\n",
        "#    ax1.fill_between(stretch, ir_t-ir_t_std*1, ir_t+ir_t_std*1, alpha=0.2, facecolor=colors(0.3), zorder=3)\n",
        "ax1.plot(stretch, ir_t, color=colors(0.3), linewidth=2, zorder=5)\n",
        "\n",
        "ax1.spines['top'].set_visible(False)\n",
        "ax1.spines['right'].set_visible(False)\n",
        "ax1.set_yticks([0.1], minor= [0.1])\n",
        "pl.title(title)\n",
        "\n",
        "legend1 = ax1.legend(loc=leg_loc,\n",
        "                     framealpha=0.85, facecolor='white')\n",
        "if 'ood' in download:\n",
        "  pl.tight_layout()\n",
        "  #pl.savefig(\"ood.pdf\", format=\"pdf\")\n",
        "  #%download_file ood.pdf\n",
        "else:\n",
        "  pl.tight_layout()\n",
        "  #pl.savefig(\"normal.pdf\", format=\"pdf\")\n",
        "  #%download_file normal.pdf\n",
        "pl.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XAbAu9iWYUVo"
      },
      "outputs": [],
      "source": [
        "#@title Gradient descent on linear regression\n",
        "number_of_datapoints = 45\n",
        "seed = 0\n",
        "\n",
        "learning_rate = 1\n",
        "num_iterations = 100\n",
        "input_size = 12\n",
        "output_size = 8\n",
        "num_data_sets = 2\n",
        "\n",
        "@partial(jax.jit, static_argnums=(1,2,3,4))\n",
        "def create_lin_reg_data(rng: jnp.ndarray, i_size, o_size, c_size, input_scale) -> dict:\n",
        "  \"\"\"Create a linear regression data set.\n",
        "     Function to learn is W*x where x ~ N(0,1) and W ~ N(0,1)\"\"\"\n",
        "  rng, new_rng, new_rng_2   = jax.random.split(rng, 3)\n",
        "  w = jax.random.normal(rng, shape=[o_size, i_size])\n",
        "  x_train = jax.random.uniform(new_rng, shape=[c_size, i_size]) *2 - 1\n",
        "  x_test = jax.random.uniform(new_rng_2, shape=[1, i_size]) *2 - 1\n",
        "  y_train = w@x_train.T\n",
        "  y_test = w@x_test.T\n",
        "  return (x_train, y_train.T, x_test, y_test.T, w)\n",
        "\n",
        "data_creator = vmap(create_lin_reg_data, in_axes=(0, None, None, None, None), out_axes=0)\n",
        "\n",
        "def l2_norm(tree):\n",
        "  \"\"\"Compute the l2 norm of a pytree of arrays. Useful for weight decay.\"\"\"\n",
        "  leaves, _ = tree_flatten(tree)\n",
        "  return 0.5*jnp.sqrt(sum(jnp.vdot(x, x) for x in leaves))\n",
        "\n",
        "def forward(W, X, Y, X_test, Y_test, lr=1., wd=0.0):\n",
        "  \"\"\"Linear regresssion.\"\"\"\n",
        "  def compute_test_loss(W):\n",
        "    P_test = jnp.einsum(\"...oj,...ij->...io\", W, X_test)\n",
        "    return 0.5*jnp.mean((Y_test-P_test)**2)\n",
        "\n",
        "  def update_step(carry, i):\n",
        "    (W, X, Y, lr, wd) = carry\n",
        "    P = jnp.einsum(\"...oj,...ij->...io\", W, X)\n",
        "    #print(W, X, P)\n",
        "    delta_W = -lr*1/X.shape[1]*(jnp.einsum(\"...ij,...ik->...jk\", P-Y, X) - wd*W)\n",
        "    return (W + delta_W, X, Y, lr, wd), (0.5*jnp.mean((Y-P)**2), compute_test_loss(W))\n",
        "  (W, _, _, _, _), losses =  jax.lax.scan(update_step, (W, X, Y, lr, wd), xs=None, length=num_iterations)\n",
        "  return W, losses\n",
        "\n",
        "def lin_reg(params, data, lr=1.) -> jnp.ndarray:\n",
        "  \"\"\"Computes the MSE loss between targets and predictions.\"\"\"\n",
        "  X, Y = data[0], data[1]\n",
        "  X_test, Y_test = data[2], data[3]\n",
        "  W, losses = forward(params, X, Y, X_test, Y_test, lr)\n",
        "  train_p = jnp.einsum(\"...oj,...ij->...io\", W, X)\n",
        "\n",
        "  test_p = jnp.einsum(\"...oj,...ij->...io\", W, X_test)\n",
        "  assert Y.shape == train_p.shape\n",
        "  assert Y_test.shape == test_p.shape\n",
        "  return (losses, jnp.mean((Y-train_p)**2), jnp.mean((Y_test-test_p)**2),\n",
        "          jnp.sqrt(jnp.sum((W-data[4])**2)),\n",
        "          train_p, test_p, W)\n",
        "\n",
        "rng = jax.random.PRNGKey(seed)\n",
        "rng, train_rng = jax.random.split(rng, 2)\n",
        "data = data_creator(jax.random.split(rng, num=num_data_sets),\n",
        "                    input_size, output_size, number_of_datapoints, 1)\n",
        "W_init = np.repeat(jax.random.normal(train_rng,\n",
        "                                     shape=[1, output_size, input_size]),\n",
        "                   num_data_sets, axis=0)\n",
        "lr_results = lin_reg(W_init, data, lr=learning_rate)\n",
        "losses_test = lr_results[0][1]\n",
        "config.key_size = input_size + output_size\n",
        "\n",
        "old_iter = config.num_layers\n",
        "config.num_layers = num_iterations"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gqtcOD75YbEa"
      },
      "outputs": [],
      "source": [
        "#@title Running the constructed Transformer forward i.e. emulating gradient descent (Note that afterwards the train_state will be overwritten and you have to start the notebook again)\n",
        "\n",
        "params_new = create_weights(input_size, output_size,\n",
        "                            number_of_datapoints, learning_rate, W_init,\n",
        "                            gd_deq=True)\n",
        "#create w_token and context\n",
        "query = jnp.concatenate([data[2], jnp.zeros_like(data[3])], axis=2)\n",
        "context = jnp.concatenate([data[0], data[1]], axis=2)\n",
        "all_data = jnp.concatenate([context, query], axis=1)\n",
        "\n",
        "losses_trans_test, losses_trans_train = [], []\n",
        "grads_trans = []\n",
        "ws_trans = []\n",
        "\n",
        "# Init query token with init prediction as y data\n",
        "# (3, 1, 5) (3, 1, 10) (3, 5, 10)\n",
        "inits = jnp.einsum(\"nkj,nij->nki\", all_data[:, -1:, :-output_size], W_init)\n",
        "\n",
        "all_data = all_data.at[:, -1:, -output_size:].set((-1)*inits)\n",
        "\n",
        "test_prediction = []\n",
        "#replace by scan\n",
        "_, _, _, rng = init()\n",
        "preds = predict.apply(params_new, rng, all_data, True)\n",
        "preds_stack = predict_stack.apply(params_new, rng, all_data, True)\n",
        "print(\"Prediction of GD: \", lr_results[-2][0])\n",
        "print(\"Prediction of trained TF: \", preds[:, -1:, -output_size:][0]*(-1))\n",
        "print(\"Differences in final predictions: \", np.linalg.norm(lr_results[-2] - preds[:, -1:, -output_size:]*(-1)))\n",
        "\n",
        "config.num_layers = old_iter"
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "CDXNLkmwhsvL"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "TPU",
    "colab": {
      "private_outputs": true,
      "provenance": [],
      "gpuType": "V28"
    },
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "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.9.6"
    },
    "vscode": {
      "interpreter": {
        "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}