gpt-2.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# GPT-2\n",
    "\n",
    "We will build a GPT-2 language model for predicting the next token in a sequence of text. The shakespeare_char dataset will be used for this demonstration, which can be found in the data folder.  \n",
    "\n",
    "GPT-2 [6] is a scaled up version of the transformer decoder. It uses multiple decoder block layers, larger embedding size, larger block size, more heads in the multi-headed attention block, and a larger vocabulary size. The vocabulary size is 50257 tokens, but for our example we will stick with the simple 65 token vocabulary. It comes in multiple sized versions, small, medium, large and XL, to accomadate for the varying compute requirements of the user. We will be creating and unofficial extra-small variant for this example due to compute constraints. If compute is not a limitation, any one of the GPT-2 variants can be created just by using the parameters shown in the **Paramaters Selection** section below. One major change from the transformer decoder notebook is the introduction of dropout layers [7]. Dropout is a regularization technique mainly used to help prevent overfitting. Seeing as we are scaling up the model, this becomes a valuable addition to the model.\n",
    "\n",
    "\n",
    "### References:\n",
    "- [1] [GPT colab notebook](https://colab.research.google.com/drive/1JMLa53HDuA-i7ZBmqV7ZnA3c_fvtXnx-?usp=sharing)\n",
    "- [2] [Video: scaling up the model!](https://www.youtube.com/watch?v=kCc8FmEb1nY&t=5869s)\n",
    "- [3] [Attention Is All You Need](https://arxiv.org/abs/1706.03762)\n",
    "- [4] [Deep Residual Learning for Image Recognition](https://arxiv.org/abs/1512.03385)\n",
    "- [5] [Layer Normalization](https://arxiv.org/abs/1607.06450)\n",
    "- [6] [Language Models are Unsupervised Multitask Learners](https://d4mucfpksywv.cloudfront.net/better-language-models/language-models.pdf)\n",
    "- [7] [Improving neural networks by preventing co-adaptation of feature detectors](https://arxiv.org/abs/1207.0580)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import requests\n",
    "import numpy as np\n",
    "import jax\n",
    "from jax import random\n",
    "import jax.numpy as jnp\n",
    "import flax.linen as nn\n",
    "import optax\n",
    "from jax import value_and_grad\n",
    "\n",
    "from helper_funcs import get_batch, generate, masked_fill, loss_fn"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parameter Selection\n",
    "\n",
    "The Parameters used below are a scaled down version of GPT-2. GPT-2 has 4 different sizes, small, medium, large and xl. This GPT-2 could be considered an extra-small version. Note that these models may not be able to fit into RAM on your device. The exact specifications of the different sized models are shown below:\n",
    "\n",
    "### GPT-2 Small\n",
    "- n_embed: 768\n",
    "- block_size: 1024\n",
    "- num_heads: 12\n",
    "- num_layers: 12\n",
    "- vocab_size: 50257 (uses Tiktoken vocab)\n",
    "\n",
    "### GPT-2 Medium\n",
    "- n_embed: 1024\n",
    "- block_size: 1024\n",
    "- num_heads: 16\n",
    "- num_layers: 24\n",
    "- vocab_size: 50257 (uses Tiktoken vocab)\n",
    "\n",
    "### GPT-2 Large\n",
    "- n_embed: 1280\n",
    "- block_size: 1024\n",
    "- num_heads: 20\n",
    "- num_layers: 36\n",
    "- vocab_size: 50257 (uses Tiktoken vocab)\n",
    "\n",
    "### GPT-2 XL\n",
    "- n_embed: 1600\n",
    "- block_size: 1024\n",
    "- num_heads: 25\n",
    "- num_layers: 48\n",
    "- vocab_size: 50257 (uses Tiktoken vocab)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_embed = 32 # Number of embedding dimensions\n",
    "batch_size = 16 # How many independent sequences will we process in parallel?\n",
    "block_size = 32 # What is the maximum context length for predictions?\n",
    "num_heads = 4 # Number of heads in the multi-headed block\n",
    "num_layers = 6 # Number of transformer decoder blocks\n",
    "drop_rate = 0.1 # Dropout rate for regularization\n",
    "\n",
    "rng_key = jax.random.PRNGKey(128)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Preparation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "length of dataset in characters: 1,115,394\n",
      "all the unique characters: \n",
      " !$&',-.3:;?ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz\n",
      "vocab size: 65\n",
      "train has 1,003,854 tokens\n",
      "val has 111,540 tokens\n"
     ]
    }
   ],
   "source": [
    "# download the tiny shakespeare dataset\n",
    "input_file_path = os.path.join('./data/shakespeare_char/input.txt')\n",
    "if not os.path.exists(input_file_path):\n",
    "    data_url = 'https://github.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt'\n",
    "    with open(input_file_path, 'w') as f:\n",
    "        f.write(requests.get(data_url).text)\n",
    "\n",
    "with open(input_file_path, 'r') as f:\n",
    "    data = f.read()\n",
    "print(f\"length of dataset in characters: {len(data):,}\")\n",
    "\n",
    "# get all the unique characters that occur in this text\n",
    "chars = sorted(list(set(data)))\n",
    "vocab_size = len(chars)\n",
    "print(\"all the unique characters:\", ''.join(chars))\n",
    "print(f\"vocab size: {vocab_size:,}\")\n",
    "\n",
    "# create a mapping from characters to integers\n",
    "stoi = { ch:i for i,ch in enumerate(chars) }\n",
    "itos = { i:ch for i,ch in enumerate(chars) }\n",
    "def encode(s):\n",
    "    return [stoi[c] for c in s] # encoder: take a string, output a list of integers\n",
    "def decode(l):\n",
    "    l = np.array(l)\n",
    "    return ''.join([itos[i] for i in l]) # decoder: take a list of integers, output a string\n",
    "\n",
    "# create the train and test splits\n",
    "n = len(data)\n",
    "train_data = data[:int(n*0.9)]\n",
    "val_data = data[int(n*0.9):]\n",
    "\n",
    "# encode both to integers\n",
    "train_ids = encode(train_data)\n",
    "val_ids = encode(val_data)\n",
    "print(f\"train has {len(train_ids):,} tokens\")\n",
    "print(f\"val has {len(val_ids):,} tokens\")\n",
    "\n",
    "train_ids = jnp.array(train_ids, dtype=jnp.uint16)\n",
    "val_ids = jnp.array(val_ids, dtype=jnp.uint16)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First Citizen:\n",
      "Before we proceed any further, hear me speak.\n",
      "\n",
      "All:\n",
      "Speak, speak.\n",
      "\n",
      "First Citizen:\n",
      "You\n"
     ]
    }
   ],
   "source": [
    "print(decode(train_ids[:100]))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build the Attention Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FeedForward(nn.Module):\n",
    "    \"\"\"\n",
    "    A feed forward multi-layer perceptron network.\n",
    "    \"\"\"\n",
    "    n_embed: int\n",
    "    drop_rate: float\n",
    "\n",
    "    @nn.compact\n",
    "    def __call__(self, x):\n",
    "        net = nn.Sequential([\n",
    "            nn.Dense(4 * self.n_embed),\n",
    "            jax.nn.relu,\n",
    "            nn.Dense(self.n_embed),\n",
    "            nn.Dropout(rate=self.drop_rate, deterministic=True)\n",
    "        ])\n",
    "        x = net(x)\n",
    "\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Head(nn.Module):\n",
    "    \"\"\"\n",
    "    A single-headed self-attention decoder block.\n",
    "    Takes the combined token and position embedding as input,\n",
    "    then calculates the key and query values.\n",
    "    The key and query are multiplied to calculate the \n",
    "    attention scores/affinities. The future weights are\n",
    "    then altered to have zero affinity, this ensures the \n",
    "    model can't \"cheat\". The input is then used to calculate\n",
    "    the values, which are then aggregated by multiplying \n",
    "    them with the weights.\n",
    "    \"\"\"\n",
    "    head_size: int\n",
    "    drop_rate: float\n",
    "\n",
    "    @nn.compact\n",
    "    def __call__(self, x):\n",
    "        B,T,C = x.shape\n",
    "        key = nn.Dense(self.head_size, use_bias=False)\n",
    "        k = key(x) # (B,T,C)\n",
    "        query = nn.Dense(self.head_size, use_bias=False)\n",
    "        q = query(x) # (B,T,C)\n",
    "        # compute attention scores (\"affinities\")\n",
    "        weights =  q @ k.transpose((0, -1, -2)) * self.head_size**-0.5 # (B, T, C) @ (B, C, T) ---> (B, T, T)\n",
    "        tril = jnp.tril(jnp.ones(shape=(T, T), dtype=bool))\n",
    "        tril = jnp.repeat(tril[None, ...], repeats=B, axis=0)\n",
    "        weights = masked_fill(tril, weights, -jnp.inf)\n",
    "        weights = jax.nn.softmax(weights, axis=-1)\n",
    "        drop = nn.Dropout(rate=self.drop_rate, deterministic=True)\n",
    "        weights = drop(weights)\n",
    "        # perform the weighted aggregation of the values\n",
    "        value = nn.Dense(self.head_size, use_bias=False)\n",
    "        v = value(x)\n",
    "        out = weights @ v\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadedAttention(nn.Module):\n",
    "    \"\"\"\n",
    "    Combines multiple heads of scaled self-attention \n",
    "    in parallel, then concatenates the heads outputs.\n",
    "    \"\"\"\n",
    "    num_heads: int\n",
    "    head_size: int\n",
    "    n_embed: int\n",
    "    drop_rate: float\n",
    "\n",
    "    @nn.compact\n",
    "    def __call__(self, x):\n",
    "        # Create a list of num_heads heads\n",
    "        heads = [Head(self.head_size, self.drop_rate) for _ in range(self.num_heads)]\n",
    "        # Provide the same input for each head\n",
    "        heads_out = [h(x) for h in heads]\n",
    "        combined_logits = jnp.concatenate(heads_out, axis=-1)\n",
    "        # Perform a linear projection of the self-attention\n",
    "        proj = nn.Dense(self.n_embed)\n",
    "        logits = proj(combined_logits)\n",
    "        drop = nn.Dropout(rate=self.drop_rate, deterministic=True)\n",
    "        logits = drop(logits)\n",
    "        return logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Block(nn.Module):\n",
    "    \"\"\"\n",
    "    Transformer decoder block.\n",
    "    It combines communication and computation.\n",
    "    The communication is performed by the \n",
    "    multi-headed attention layer.\n",
    "    Then the computation is performed by \n",
    "    the feed forward block.\n",
    "    Skip connections are used to make the block scalable \n",
    "    and layer norm is used to speed up training.\n",
    "    \"\"\"\n",
    "    n_embed: int\n",
    "    num_heads: int\n",
    "    drop_rate: float\n",
    "\n",
    "    @nn.compact\n",
    "    def __call__(self, x):\n",
    "        head_size = self.n_embed // self.num_heads\n",
    "        sa_heads = MultiHeadedAttention(self.num_heads, head_size, self.n_embed, self.drop_rate)\n",
    "        # Using skip connections with x + heads\n",
    "        x = x + sa_heads(nn.LayerNorm()(x)) # apply one head of self-attention (B, T, C)\n",
    "        ffwd = FeedForward(self.n_embed, self.drop_rate)\n",
    "        x = x + ffwd(nn.LayerNorm()(x))\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GPT2(nn.Module):\n",
    "    \"\"\"\n",
    "    GPT-2 language model.\n",
    "    Uses the previous tokens in the sequence to \n",
    "    determine the probabilities of the next token.\n",
    "    Processes the combined position and token embedding\n",
    "    through multiple transformer decoder blocks, \n",
    "    which is then processed through a dense layer to \n",
    "    aquire the token logits.\n",
    "    The logits can then be processed through a softmax\n",
    "    function to calculate the token probabilities.\n",
    "    \"\"\"\n",
    "    vocab_size: int\n",
    "    n_embed: int\n",
    "    block_size: int\n",
    "    num_heads: int\n",
    "    num_layers: int\n",
    "    drop_rate: float\n",
    "    \n",
    "    @nn.compact\n",
    "    def __call__(self, index_seq):\n",
    "        B, T = index_seq.shape\n",
    "\n",
    "        # Each token directly reads off the logits for the next token from a lookup table\n",
    "        token_embedding_table = nn.Embed(num_embeddings=self.vocab_size, features=self.n_embed) \n",
    "        token_emb = token_embedding_table(index_seq) # (B, T, C)\n",
    "\n",
    "        position_embedding_table = nn.Embed(num_embeddings=self.block_size, features=self.n_embed) \n",
    "        pos_emb = position_embedding_table(jnp.arange(T)) # (T, C)\n",
    "\n",
    "        x = token_emb + pos_emb # (B, T, C)\n",
    "\n",
    "        decoder_blocks = [Block(self.n_embed, num_heads=self.num_heads, drop_rate=self.drop_rate) for _ in range(self.num_layers)]\n",
    "        decoder_blocks.append(nn.LayerNorm())\n",
    "        blocks = nn.Sequential(\n",
    "            decoder_blocks\n",
    "        )\n",
    "        x = blocks(x)\n",
    "\n",
    "        lm_head = nn.Dense(self.vocab_size)\n",
    "        logits = lm_head(x) # (B, T, vocab_size)\n",
    "\n",
    "        return logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = GPT2(vocab_size, n_embed, block_size, num_heads, num_layers, drop_rate)\n",
    "dummy_x = jnp.zeros(shape=(batch_size, block_size), dtype=jnp.uint16)\n",
    "variables = model.init(rng_key, dummy_x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1, 32, 65)\n"
     ]
    }
   ],
   "source": [
    "out = model.apply(variables, dummy_x)\n",
    "print(out.shape)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Text Generation Pre-Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated text: \n",
      "\n",
      "EiQnlVvj''mtKvgrI w:\n"
     ]
    }
   ],
   "source": [
    "index_seq = jnp.zeros(shape=(1,1), dtype=jnp.uint16)\n",
    "max_new_tokens = 20\n",
    "\n",
    "generated_indices = generate(variables, model.apply, index_seq, rng_key, vocab_size, 1, block_size, max_new_tokens)\n",
    "generated_indices = list(np.array(generated_indices[0]))\n",
    "print(\"Generated text: \")\n",
    "print(decode(generated_indices))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train the Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = optax.adamw(learning_rate=1e-2)\n",
    "opt_state = optimizer.init(variables)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 0, Loss: 4.4984\n",
      "Epoch: 1, Loss: 4.0247\n",
      "Epoch: 2, Loss: 3.7785\n",
      "Epoch: 3, Loss: 3.4819\n",
      "Epoch: 4, Loss: 3.5813\n",
      "Epoch: 5, Loss: 3.4390\n",
      "Epoch: 6, Loss: 3.3522\n",
      "Epoch: 7, Loss: 3.3741\n",
      "Epoch: 8, Loss: 3.2018\n",
      "Epoch: 9, Loss: 3.5234\n",
      "Epoch: 10, Loss: 3.2206\n",
      "Epoch: 11, Loss: 3.1798\n",
      "Epoch: 12, Loss: 3.2482\n",
      "Epoch: 13, Loss: 3.1972\n",
      "Epoch: 14, Loss: 3.2870\n",
      "Epoch: 15, Loss: 3.2179\n",
      "Epoch: 16, Loss: 3.2862\n",
      "Epoch: 17, Loss: 3.3190\n",
      "Epoch: 18, Loss: 3.3815\n",
      "Epoch: 19, Loss: 3.3236\n",
      "Epoch: 20, Loss: 3.2920\n",
      "Epoch: 21, Loss: 3.1990\n",
      "Epoch: 22, Loss: 3.3100\n",
      "Epoch: 23, Loss: 3.2309\n",
      "Epoch: 24, Loss: 3.3584\n",
      "Epoch: 25, Loss: 3.2649\n",
      "Epoch: 26, Loss: 3.1731\n",
      "Epoch: 27, Loss: 3.1962\n",
      "Epoch: 28, Loss: 3.4236\n",
      "Epoch: 29, Loss: 3.1557\n",
      "Epoch: 30, Loss: 3.2859\n",
      "Epoch: 31, Loss: 3.1309\n",
      "Epoch: 32, Loss: 3.1754\n",
      "Epoch: 33, Loss: 3.1635\n",
      "Epoch: 34, Loss: 3.2867\n",
      "Epoch: 35, Loss: 3.1182\n",
      "Epoch: 36, Loss: 2.9672\n",
      "Epoch: 37, Loss: 3.1443\n",
      "Epoch: 38, Loss: 3.1414\n",
      "Epoch: 39, Loss: 3.1504\n",
      "Epoch: 40, Loss: 3.0108\n",
      "Epoch: 41, Loss: 3.0403\n",
      "Epoch: 42, Loss: 2.9964\n",
      "Epoch: 43, Loss: 3.0329\n",
      "Epoch: 44, Loss: 3.1116\n",
      "Epoch: 45, Loss: 2.9713\n",
      "Epoch: 46, Loss: 2.9396\n",
      "Epoch: 47, Loss: 3.0610\n",
      "Epoch: 48, Loss: 3.0597\n",
      "Epoch: 49, Loss: 2.8929\n",
      "Epoch: 50, Loss: 2.8312\n",
      "Epoch: 51, Loss: 2.8679\n",
      "Epoch: 52, Loss: 3.0265\n",
      "Epoch: 53, Loss: 2.9092\n",
      "Epoch: 54, Loss: 2.8468\n",
      "Epoch: 55, Loss: 2.8166\n",
      "Epoch: 56, Loss: 2.9492\n",
      "Epoch: 57, Loss: 2.9896\n",
      "Epoch: 58, Loss: 2.9172\n",
      "Epoch: 59, Loss: 2.8209\n",
      "Epoch: 60, Loss: 2.8746\n",
      "Epoch: 61, Loss: 2.8864\n",
      "Epoch: 62, Loss: 2.7877\n",
      "Epoch: 63, Loss: 2.8333\n",
      "Epoch: 64, Loss: 2.8097\n",
      "Epoch: 65, Loss: 2.9920\n",
      "Epoch: 66, Loss: 2.8437\n",
      "Epoch: 67, Loss: 2.8001\n",
      "Epoch: 68, Loss: 2.6927\n",
      "Epoch: 69, Loss: 2.8162\n",
      "Epoch: 70, Loss: 2.8116\n",
      "Epoch: 71, Loss: 2.7609\n",
      "Epoch: 72, Loss: 2.8246\n",
      "Epoch: 73, Loss: 2.7641\n",
      "Epoch: 74, Loss: 2.6625\n",
      "Epoch: 75, Loss: 2.7137\n",
      "Epoch: 76, Loss: 2.7618\n",
      "Epoch: 77, Loss: 2.7533\n",
      "Epoch: 78, Loss: 2.6671\n",
      "Epoch: 79, Loss: 2.9097\n",
      "Epoch: 80, Loss: 2.7236\n",
      "Epoch: 81, Loss: 2.6633\n",
      "Epoch: 82, Loss: 2.6982\n",
      "Epoch: 83, Loss: 2.7677\n",
      "Epoch: 84, Loss: 2.7462\n",
      "Epoch: 85, Loss: 2.6724\n",
      "Epoch: 86, Loss: 2.6885\n",
      "Epoch: 87, Loss: 2.6859\n",
      "Epoch: 88, Loss: 2.7013\n",
      "Epoch: 89, Loss: 2.8477\n",
      "Epoch: 90, Loss: 2.7398\n",
      "Epoch: 91, Loss: 2.6856\n",
      "Epoch: 92, Loss: 2.6711\n",
      "Epoch: 93, Loss: 2.6841\n",
      "Epoch: 94, Loss: 2.7377\n",
      "Epoch: 95, Loss: 2.6989\n",
      "Epoch: 96, Loss: 2.8792\n",
      "Epoch: 97, Loss: 2.6758\n",
      "Epoch: 98, Loss: 2.7484\n",
      "Epoch: 99, Loss: 2.7810\n"
     ]
    }
   ],
   "source": [
    "steps = 100\n",
    "\n",
    "for step in range(steps):\n",
    "    rng_key, subkey = jax.random.split(rng_key)\n",
    "    xb, yb = get_batch(train_ids, subkey, batch_size, block_size)\n",
    "\n",
    "    loss, grads = value_and_grad(loss_fn, argnums=(0))(\n",
    "        variables, \n",
    "        model.apply,\n",
    "        xb, \n",
    "        yb\n",
    "    )\n",
    "    updates, opt_state = optimizer.update(grads, opt_state, variables)\n",
    "    variables = optax.apply_updates(variables, updates) \n",
    "\n",
    "    print(f\"Epoch: {step}, Loss: {loss :.4f}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Text Generation Post-Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated text: \n",
      "\n",
      "Wos at arXOTerof t t\n"
     ]
    }
   ],
   "source": [
    "index_seq = jnp.zeros(shape=(1,1), dtype=jnp.uint16)\n",
    "max_new_tokens = 20\n",
    "\n",
    "generated_indices = generate(variables, model.apply, index_seq, rng_key, vocab_size, 1, block_size, max_new_tokens)\n",
    "generated_indices = list(np.array(generated_indices[0]))\n",
    "print(\"Generated text: \")\n",
    "print(decode(generated_indices))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "flax2",
   "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.16"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "61a659281f728d55fc033691f61c254a0324a8156d8c726d9ae327f7f6499663"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}