70b_sample.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "!sudo rm /tmp/libtpu_lockfile\n",
    "import penzai\n",
    "from penzai import pz\n",
    "import jax.numpy as jnp\n",
    "import jax_smi\n",
    "import jax\n",
    "jax_smi.initialise_tracking()\n",
    "pz.ts.register_as_default()\n",
    "pz.ts.register_autovisualize_magic()\n",
    "pz.enable_interactive_context()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_seq_len = 64\n",
    "batch_size = 16"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Folding and caching...\n"
     ]
    },
    {
     "ename": "TypeError",
     "evalue": "can't multiply sequence by non-int of type 'jaxlib.xla_extension.PjitFunction'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[3], line 11\u001b[0m\n\u001b[1;32m      5\u001b[0m llama \u001b[38;5;241m=\u001b[39m LlamaTransformer\u001b[38;5;241m.\u001b[39mfrom_pretrained((\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmodels/llama-3-70b-1.gguf\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmodels/llama-3-70b-2.gguf\u001b[39m\u001b[38;5;124m\"\u001b[39m),\n\u001b[1;32m      6\u001b[0m                                          device_map\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mauto:mp=4\u001b[39m\u001b[38;5;124m\"\u001b[39m, load_on_cpu\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m,\n\u001b[1;32m      7\u001b[0m                                         \u001b[38;5;66;03m#  transpose_rotary=False\u001b[39;00m\n\u001b[1;32m      8\u001b[0m                                         \u001b[38;5;66;03m#  transpose_rotary=True\u001b[39;00m\n\u001b[1;32m      9\u001b[0m                                          )\n\u001b[1;32m     10\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mFolding and caching...\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m---> 11\u001b[0m llama, cache \u001b[38;5;241m=\u001b[39m \u001b[43mFoldedLlamaKVCachingTransformer\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfrom_uncached\u001b[49m\u001b[43m(\u001b[49m\u001b[43mllama\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmax_seq_len\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m{\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mbatch\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mbatch_size\u001b[49m\u001b[43m}\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     12\u001b[0m \u001b[38;5;66;03m# llama = flashify(llama)\u001b[39;00m\n",
      "File \u001b[0;32m~/micrlhf-progress/micrlhf/caching_llama.py:249\u001b[0m, in \u001b[0;36mFoldedLlamaKVCachingTransformer.from_uncached\u001b[0;34m(cls, uncached, cache_len, batch_axes, fold_kwargs)\u001b[0m\n\u001b[1;32m    247\u001b[0m     fold_kwargs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mdict\u001b[39m()\n\u001b[1;32m    248\u001b[0m fold_kwargs \u001b[38;5;241m=\u001b[39m {\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mreturn_aux\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;28;01mTrue\u001b[39;00m, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mfold_kwargs}\n\u001b[0;32m--> 249\u001b[0m uncached, aux \u001b[38;5;241m=\u001b[39m \u001b[43msequential_to_scan\u001b[49m\u001b[43m(\u001b[49m\u001b[43muncached\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mfold_kwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    250\u001b[0m n_layers \u001b[38;5;241m=\u001b[39m aux[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mn_layers\u001b[39m\u001b[38;5;124m\"\u001b[39m]\n\u001b[1;32m    251\u001b[0m cached_axes \u001b[38;5;241m=\u001b[39m {\n\u001b[1;32m    252\u001b[0m         \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mbatch_axes,\n\u001b[1;32m    253\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mprojection\u001b[39m\u001b[38;5;124m\"\u001b[39m: uncached\u001b[38;5;241m.\u001b[39mconfig\u001b[38;5;241m.\u001b[39mprojection_dim,\n\u001b[1;32m    254\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mlayer\u001b[39m\u001b[38;5;124m\"\u001b[39m: n_layers,\n\u001b[1;32m    255\u001b[0m }\n",
      "File \u001b[0;32m~/micrlhf-progress/micrlhf/scan.py:74\u001b[0m, in \u001b[0;36msequential_to_scan\u001b[0;34m(model, sequential_n, return_aux, save_to_cpu)\u001b[0m\n\u001b[1;32m     71\u001b[0m     aux[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mn_layers\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mlen\u001b[39m(layers)\n\u001b[1;32m     72\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m folded\n\u001b[0;32m---> 74\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[43mpick_sequential\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mselect\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msequential_n\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mapply\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfn\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     75\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m return_aux:\n\u001b[1;32m     76\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m result, aux\n",
      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/micrlhf-progress-a058ydGG-py3.12/lib/python3.12/site-packages/penzai/core/selectors.py:322\u001b[0m, in \u001b[0;36mSelection.apply\u001b[0;34m(self, fn, with_keypath, keep_selected)\u001b[0m\n\u001b[1;32m    317\u001b[0m   new_values \u001b[38;5;241m=\u001b[39m collections\u001b[38;5;241m.\u001b[39mOrderedDict(\n\u001b[1;32m    318\u001b[0m       [(k, fn(k, v)) \u001b[38;5;28;01mfor\u001b[39;00m k, v \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mselected_by_path\u001b[38;5;241m.\u001b[39mitems()]\n\u001b[1;32m    319\u001b[0m   )\n\u001b[1;32m    320\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    321\u001b[0m   new_values \u001b[38;5;241m=\u001b[39m collections\u001b[38;5;241m.\u001b[39mOrderedDict(\n\u001b[0;32m--> 322\u001b[0m       [(k, \u001b[43mfn\u001b[49m\u001b[43m(\u001b[49m\u001b[43mv\u001b[49m\u001b[43m)\u001b[49m) \u001b[38;5;28;01mfor\u001b[39;00m k, v \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mselected_by_path\u001b[38;5;241m.\u001b[39mitems()]\n\u001b[1;32m    323\u001b[0m   )\n\u001b[1;32m    324\u001b[0m new_selection \u001b[38;5;241m=\u001b[39m Selection(\n\u001b[1;32m    325\u001b[0m     selected_by_path\u001b[38;5;241m=\u001b[39mnew_values, remainder\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mremainder\n\u001b[1;32m    326\u001b[0m )\n\u001b[1;32m    327\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m keep_selected:\n",
      "File \u001b[0;32m~/micrlhf-progress/micrlhf/scan.py:63\u001b[0m, in \u001b[0;36msequential_to_scan.<locals>.fn\u001b[0;34m(seq)\u001b[0m\n\u001b[1;32m     58\u001b[0m     w \u001b[38;5;241m=\u001b[39m jax\u001b[38;5;241m.\u001b[39mtree\u001b[38;5;241m.\u001b[39mmap(\u001b[38;5;28;01mlambda\u001b[39;00m \u001b[38;5;241m*\u001b[39mws: pz\u001b[38;5;241m.\u001b[39mnx\u001b[38;5;241m.\u001b[39mNamedArray(\n\u001b[1;32m     59\u001b[0m         data_array\u001b[38;5;241m=\u001b[39mnp\u001b[38;5;241m.\u001b[39marray([w\u001b[38;5;241m.\u001b[39mdata_array \u001b[38;5;28;01mfor\u001b[39;00m w \u001b[38;5;129;01min\u001b[39;00m ws]),\n\u001b[1;32m     60\u001b[0m         named_axes\u001b[38;5;241m=\u001b[39mOrderedDict([(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mlayer\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;28mlen\u001b[39m(ws))] \u001b[38;5;241m+\u001b[39m \u001b[38;5;28mlist\u001b[39m(ws[\u001b[38;5;241m0\u001b[39m]\u001b[38;5;241m.\u001b[39mnamed_axes\u001b[38;5;241m.\u001b[39mitems()))\n\u001b[1;32m     61\u001b[0m     ), \u001b[38;5;241m*\u001b[39mws, is_leaf\u001b[38;5;241m=\u001b[39mis_nx)\n\u001b[1;32m     62\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m---> 63\u001b[0m     w \u001b[38;5;241m=\u001b[39m jax\u001b[38;5;241m.\u001b[39mtree\u001b[38;5;241m.\u001b[39mmap(\u001b[43mjax\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mjit\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43;01mlambda\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mws\u001b[49m\u001b[43m:\u001b[49m\n\u001b[1;32m     64\u001b[0m \u001b[43m        \u001b[49m\u001b[43mpz\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mnx\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstack\u001b[49m\u001b[43m(\u001b[49m\u001b[43mws\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mlayer\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43muntag\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mlayer\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mwith_positional_prefix\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtag\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mlayer\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     65\u001b[0m \u001b[43m                     \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mws\u001b[49m, is_leaf\u001b[38;5;241m=\u001b[39mis_nx)\n\u001b[1;32m     67\u001b[0m weight \u001b[38;5;241m=\u001b[39m jax\u001b[38;5;241m.\u001b[39mtree_util\u001b[38;5;241m.\u001b[39mtree_unflatten(tds[\u001b[38;5;241m0\u001b[39m], w)\n\u001b[1;32m     68\u001b[0m layer \u001b[38;5;241m=\u001b[39m eqx\u001b[38;5;241m.\u001b[39mcombine(weight, treedefs[\u001b[38;5;241m0\u001b[39m], is_leaf\u001b[38;5;241m=\u001b[39mis_nx)\n",
      "\u001b[0;31mTypeError\u001b[0m: can't multiply sequence by non-int of type 'jaxlib.xla_extension.PjitFunction'"
     ]
    }
   ],
   "source": [
    "from micrlhf.llama import LlamaTransformer\n",
    "from micrlhf.caching_llama import FoldedLlamaKVCachingTransformer\n",
    "from micrlhf.flash import flashify\n",
    "\n",
    "llama = LlamaTransformer.from_pretrained((\"models/llama-3-70b-1.gguf\", \"models/llama-3-70b-2.gguf\"),\n",
    "                                         device_map=\"auto:mp=4\", load_on_cpu=True,\n",
    "                                        #  transpose_rotary=False\n",
    "                                        #  transpose_rotary=True\n",
    "                                         )\n",
    "print(\"Folding and caching...\")\n",
    "llama, cache = FoldedLlamaKVCachingTransformer.from_uncached(llama, max_seq_len, {\"batch\": batch_size})\n",
    "# llama = flashify(llama)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "llama = llama.to_tpu()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gc\n",
    "gc.collect();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/neverix/.cache/pypoetry/virtualenvs/micrlhf-progress-a058ydGG-py3.12/lib/python3.12/site-packages/huggingface_hub/file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
      "  warnings.warn(\n",
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoTokenizer\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"nvidia/Llama3-ChatQA-1.5-70B\")\n",
    "prompt = tokenizer.apply_chat_template([\n",
    "    {\"role\": \"system\", \"content\": \"You are a pirate chatbot who always responds in pirate speak!\"},\n",
    "    {\"role\": \"user\", \"content\": \"Who are you?\"},\n",
    "], tokenize=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "ename": "StructureMismatchError",
     "evalue": "Error while checking the input of a layer of type NamedEinsum:\nMismatch while checking structures:\nAt root[1]: Named shape mismatch between value {'batch': 16, 'projection': 128, 'layer': 80, 'kv_heads': 8, 'seq': 64} and pattern {'seq': var('I')['tkv'], 'kv_heads': var('I')['h'], 'projection': var('I')['p'], 'batch': var('B')['batch']:=16}:\n  Unexpected names in value's named shape: ['layer']\n  After inlining var('B') = {'batch': 16} from root[0]",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mStructureMismatchError\u001b[0m                    Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[9], line 2\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmicrlhf\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01msampling\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m sample\n\u001b[0;32m----> 2\u001b[0m texts, cache \u001b[38;5;241m=\u001b[39m \u001b[43msample\u001b[49m\u001b[43m(\u001b[49m\u001b[43m(\u001b[49m\u001b[43mllama\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcache\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtokenizer\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mprompt\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbatch_size\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mbatch_size\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmax_seq_len\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mmax_seq_len\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m      3\u001b[0m \u001b[43m       \u001b[49m\u001b[43mdo_sample\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mreturn_only_completion\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m)\u001b[49m\n\u001b[1;32m      4\u001b[0m texts\n",
      "File \u001b[0;32m~/micrlhf-progress/micrlhf/sampling.py:111\u001b[0m, in \u001b[0;36msample\u001b[0;34m(llama, tokenizer, prompt, batch_size, max_seq_len, pad_token_id, do_sample, return_model, strip_padding, return_only_completion, seed, verbose, use_jit, cache_kwargs)\u001b[0m\n\u001b[1;32m    106\u001b[0m base_inputs \u001b[38;5;241m=\u001b[39m dataclasses\u001b[38;5;241m.\u001b[39mreplace(base_inputs,\n\u001b[1;32m    107\u001b[0m                                   positions\u001b[38;5;241m=\u001b[39mbase_inputs\u001b[38;5;241m.\u001b[39mpositions \u001b[38;5;241m-\u001b[39m offsets,\n\u001b[1;32m    108\u001b[0m                                   attention_mask\u001b[38;5;241m=\u001b[39mbase_inputs\u001b[38;5;241m.\u001b[39mattention_mask \u001b[38;5;241m&\u001b[39m base_mask\u001b[38;5;241m.\u001b[39muntag(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mseq\u001b[39m\u001b[38;5;124m\"\u001b[39m)\u001b[38;5;241m.\u001b[39mtag(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mkv_seq\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m    109\u001b[0m                                   )\n\u001b[1;32m    110\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m jax\u001b[38;5;241m.\u001b[39mdisable_jit(disable\u001b[38;5;241m=\u001b[39m\u001b[38;5;129;01mnot\u001b[39;00m use_jit):\n\u001b[0;32m--> 111\u001b[0m     logits, cache \u001b[38;5;241m=\u001b[39m \u001b[43mcall_llama\u001b[49m\u001b[43m(\u001b[49m\u001b[43mllama_cached\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbase_inputs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    112\u001b[0m     cache \u001b[38;5;241m=\u001b[39m dataclasses\u001b[38;5;241m.\u001b[39mreplace(cache, cache_end_index\u001b[38;5;241m=\u001b[39minitial_length)\n\u001b[1;32m    114\u001b[0m     key \u001b[38;5;241m=\u001b[39m jax\u001b[38;5;241m.\u001b[39mrandom\u001b[38;5;241m.\u001b[39mkey(seed \u001b[38;5;28;01mif\u001b[39;00m seed \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;28;01melse\u001b[39;00m random\u001b[38;5;241m.\u001b[39mrandrange(\u001b[38;5;241m0\u001b[39m, \u001b[38;5;241m2\u001b[39m\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m32\u001b[39m))\n",
      "    \u001b[0;31m[... skipping hidden 11 frame]\u001b[0m\n",
      "File \u001b[0;32m~/micrlhf-progress/micrlhf/sampling.py:24\u001b[0m, in \u001b[0;36mcall_llama\u001b[0;34m(llama_cached, inputs)\u001b[0m\n\u001b[1;32m     22\u001b[0m \u001b[38;5;129m@jax\u001b[39m\u001b[38;5;241m.\u001b[39mjit\n\u001b[1;32m     23\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mcall_llama\u001b[39m(llama_cached, inputs):\n\u001b[0;32m---> 24\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mllama_cached\u001b[49m\u001b[43m(\u001b[49m\u001b[43minputs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/micrlhf-progress-a058ydGG-py3.12/lib/python3.12/site-packages/penzai/core/layer.py:82\u001b[0m, in \u001b[0;36mchecked_layer_call.<locals>.wrapper\u001b[0;34m(self, argument)\u001b[0m\n\u001b[1;32m     80\u001b[0m output_structure \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39moutput_structure()\n\u001b[1;32m     81\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m jax\u001b[38;5;241m.\u001b[39mnamed_scope(\u001b[38;5;28mtype\u001b[39m(\u001b[38;5;28mself\u001b[39m)\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m):\n\u001b[0;32m---> 82\u001b[0m   result \u001b[38;5;241m=\u001b[39m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margument\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     83\u001b[0m shapecheck\u001b[38;5;241m.\u001b[39mcheck_structure(\n\u001b[1;32m     84\u001b[0m     result,\n\u001b[1;32m     85\u001b[0m     pattern\u001b[38;5;241m=\u001b[39moutput_structure,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     90\u001b[0m     ),\n\u001b[1;32m     91\u001b[0m )\n\u001b[1;32m     92\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m result\n",
      "File \u001b[0;32m~/micrlhf-progress/micrlhf/caching_llama.py:219\u001b[0m, in \u001b[0;36mFoldedLlamaKVCachingTransformer.__call__\u001b[0;34m(self, inputs)\u001b[0m\n\u001b[1;32m    217\u001b[0m \u001b[38;5;129m@pz\u001b[39m\u001b[38;5;241m.\u001b[39mchecked_layer_call\n\u001b[1;32m    218\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__call__\u001b[39m(\u001b[38;5;28mself\u001b[39m, inputs: LlamaKVCachingInputs) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m \u001b[38;5;28mtuple\u001b[39m[pz\u001b[38;5;241m.\u001b[39mnx\u001b[38;5;241m.\u001b[39mNamedArray, LlamaKVCachingState]:\n\u001b[0;32m--> 219\u001b[0m     outs, kv_caches \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbody\u001b[49m\u001b[43m(\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    220\u001b[0m \u001b[43m            \u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    221\u001b[0m \u001b[43m                    \u001b[49m\u001b[43mLlamaInputs\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtokens\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minputs\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtokens\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    222\u001b[0m \u001b[43m                                \u001b[49m\u001b[43mpositions\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minputs\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mpositions\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    223\u001b[0m \u001b[43m                                \u001b[49m\u001b[43mattention_mask\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minputs\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mattention_mask\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    224\u001b[0m \u001b[43m                    \u001b[49m\u001b[43minputs\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msampling_state\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcache_end_index\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    225\u001b[0m \u001b[43m            \u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    226\u001b[0m \u001b[43m            \u001b[49m\u001b[43minputs\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msampling_state\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mkv_caches\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    227\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    228\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m outs, LlamaKVCachingState(\n\u001b[1;32m    229\u001b[0m             cache_len\u001b[38;5;241m=\u001b[39minputs\u001b[38;5;241m.\u001b[39msampling_state\u001b[38;5;241m.\u001b[39mcache_len,\n\u001b[1;32m    230\u001b[0m             batch_axes\u001b[38;5;241m=\u001b[39minputs\u001b[38;5;241m.\u001b[39msampling_state\u001b[38;5;241m.\u001b[39mbatch_axes,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    235\u001b[0m             ),\n\u001b[1;32m    236\u001b[0m     )\n",
      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/micrlhf-progress-a058ydGG-py3.12/lib/python3.12/site-packages/penzai/core/layer.py:82\u001b[0m, in \u001b[0;36mchecked_layer_call.<locals>.wrapper\u001b[0;34m(self, argument)\u001b[0m\n\u001b[1;32m     80\u001b[0m output_structure \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39moutput_structure()\n\u001b[1;32m     81\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m jax\u001b[38;5;241m.\u001b[39mnamed_scope(\u001b[38;5;28mtype\u001b[39m(\u001b[38;5;28mself\u001b[39m)\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m):\n\u001b[0;32m---> 82\u001b[0m   result \u001b[38;5;241m=\u001b[39m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margument\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     83\u001b[0m shapecheck\u001b[38;5;241m.\u001b[39mcheck_structure(\n\u001b[1;32m     84\u001b[0m     result,\n\u001b[1;32m     85\u001b[0m     pattern\u001b[38;5;241m=\u001b[39moutput_structure,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     90\u001b[0m     ),\n\u001b[1;32m     91\u001b[0m )\n\u001b[1;32m     92\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m result\n",
      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/micrlhf-progress-a058ydGG-py3.12/lib/python3.12/site-packages/penzai/data_effects/local_state.py:298\u001b[0m, in \u001b[0;36mWithFunctionalLocalState.__call__\u001b[0;34m(self, argument)\u001b[0m\n\u001b[1;32m    288\u001b[0m impls \u001b[38;5;241m=\u001b[39m {\n\u001b[1;32m    289\u001b[0m     k: LocalStateEffectImpl(_state\u001b[38;5;241m=\u001b[39mv, _handler_id\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mhandler_id)\n\u001b[1;32m    290\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m k, v \u001b[38;5;129;01min\u001b[39;00m states\u001b[38;5;241m.\u001b[39mitems()\n\u001b[1;32m    291\u001b[0m }\n\u001b[1;32m    292\u001b[0m handled_body \u001b[38;5;241m=\u001b[39m (\n\u001b[1;32m    293\u001b[0m     selectors\u001b[38;5;241m.\u001b[39mselect(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mbody)\n\u001b[1;32m    294\u001b[0m     \u001b[38;5;241m.\u001b[39mat_instances_of(HandledLocalStateRef)\n\u001b[1;32m    295\u001b[0m     \u001b[38;5;241m.\u001b[39mwhere(\u001b[38;5;28;01mlambda\u001b[39;00m ref: ref\u001b[38;5;241m.\u001b[39mhandler_id \u001b[38;5;241m==\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mhandler_id)\n\u001b[1;32m    296\u001b[0m     \u001b[38;5;241m.\u001b[39mapply(\u001b[38;5;28;01mlambda\u001b[39;00m ref: impls[ref\u001b[38;5;241m.\u001b[39mname])\n\u001b[1;32m    297\u001b[0m )\n\u001b[0;32m--> 298\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[43mhandled_body\u001b[49m\u001b[43m(\u001b[49m\u001b[43minner_arg\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    299\u001b[0m new_states \u001b[38;5;241m=\u001b[39m {k: impl\u001b[38;5;241m.\u001b[39m_state \u001b[38;5;28;01mfor\u001b[39;00m k, impl \u001b[38;5;129;01min\u001b[39;00m impls\u001b[38;5;241m.\u001b[39mitems()}\n\u001b[1;32m    300\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m result, new_states\n",
      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/micrlhf-progress-a058ydGG-py3.12/lib/python3.12/site-packages/penzai/core/layer.py:82\u001b[0m, in \u001b[0;36mchecked_layer_call.<locals>.wrapper\u001b[0;34m(self, argument)\u001b[0m\n\u001b[1;32m     80\u001b[0m output_structure \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39moutput_structure()\n\u001b[1;32m     81\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m jax\u001b[38;5;241m.\u001b[39mnamed_scope(\u001b[38;5;28mtype\u001b[39m(\u001b[38;5;28mself\u001b[39m)\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m):\n\u001b[0;32m---> 82\u001b[0m   result \u001b[38;5;241m=\u001b[39m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margument\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     83\u001b[0m shapecheck\u001b[38;5;241m.\u001b[39mcheck_structure(\n\u001b[1;32m     84\u001b[0m     result,\n\u001b[1;32m     85\u001b[0m     pattern\u001b[38;5;241m=\u001b[39moutput_structure,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     90\u001b[0m     ),\n\u001b[1;32m     91\u001b[0m )\n\u001b[1;32m     92\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m result\n",
      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/micrlhf-progress-a058ydGG-py3.12/lib/python3.12/site-packages/penzai/data_effects/side_input.py:146\u001b[0m, in \u001b[0;36mWithSideInputsFromInputTuple.__call__\u001b[0;34m(self, argument)\u001b[0m\n\u001b[1;32m    136\u001b[0m impls \u001b[38;5;241m=\u001b[39m {\n\u001b[1;32m    137\u001b[0m     tag: SideInputEffectImpl(_value\u001b[38;5;241m=\u001b[39mval, _handler_id\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mhandler_id)\n\u001b[1;32m    138\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m tag, val \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mzip\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mside_input_tags, side_inputs)\n\u001b[1;32m    139\u001b[0m }\n\u001b[1;32m    140\u001b[0m handled_body \u001b[38;5;241m=\u001b[39m (\n\u001b[1;32m    141\u001b[0m     selectors\u001b[38;5;241m.\u001b[39mselect(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mbody)\n\u001b[1;32m    142\u001b[0m     \u001b[38;5;241m.\u001b[39mat_instances_of(HandledSideInputRef)\n\u001b[1;32m    143\u001b[0m     \u001b[38;5;241m.\u001b[39mwhere(\u001b[38;5;28;01mlambda\u001b[39;00m ref: ref\u001b[38;5;241m.\u001b[39mhandler_id \u001b[38;5;241m==\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mhandler_id)\n\u001b[1;32m    144\u001b[0m     \u001b[38;5;241m.\u001b[39mapply(\u001b[38;5;28;01mlambda\u001b[39;00m ref: impls[ref\u001b[38;5;241m.\u001b[39mtag])\n\u001b[1;32m    145\u001b[0m )\n\u001b[0;32m--> 146\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mhandled_body\u001b[49m\u001b[43m(\u001b[49m\u001b[43minner_arg\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/micrlhf-progress-a058ydGG-py3.12/lib/python3.12/site-packages/penzai/core/layer.py:82\u001b[0m, in \u001b[0;36mchecked_layer_call.<locals>.wrapper\u001b[0;34m(self, argument)\u001b[0m\n\u001b[1;32m     80\u001b[0m output_structure \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39moutput_structure()\n\u001b[1;32m     81\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m jax\u001b[38;5;241m.\u001b[39mnamed_scope(\u001b[38;5;28mtype\u001b[39m(\u001b[38;5;28mself\u001b[39m)\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m):\n\u001b[0;32m---> 82\u001b[0m   result \u001b[38;5;241m=\u001b[39m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margument\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     83\u001b[0m shapecheck\u001b[38;5;241m.\u001b[39mcheck_structure(\n\u001b[1;32m     84\u001b[0m     result,\n\u001b[1;32m     85\u001b[0m     pattern\u001b[38;5;241m=\u001b[39moutput_structure,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     90\u001b[0m     ),\n\u001b[1;32m     91\u001b[0m )\n\u001b[1;32m     92\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m result\n",
      "File \u001b[0;32m~/micrlhf-progress/micrlhf/llama.py:256\u001b[0m, in \u001b[0;36mLlamaTransformer.__call__\u001b[0;34m(self, inputs)\u001b[0m\n\u001b[1;32m    254\u001b[0m \u001b[38;5;129m@pz\u001b[39m\u001b[38;5;241m.\u001b[39mchecked_layer_call\n\u001b[1;32m    255\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__call__\u001b[39m(\u001b[38;5;28mself\u001b[39m, inputs: LlamaInputs) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m pz\u001b[38;5;241m.\u001b[39mnx\u001b[38;5;241m.\u001b[39mNamedArray:\n\u001b[0;32m--> 256\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbody\u001b[49m\u001b[43m(\u001b[49m\u001b[43m(\u001b[49m\u001b[43minputs\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtokens\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minputs\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mpositions\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minputs\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mattention_mask\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/micrlhf-progress-a058ydGG-py3.12/lib/python3.12/site-packages/penzai/core/layer.py:82\u001b[0m, in \u001b[0;36mchecked_layer_call.<locals>.wrapper\u001b[0;34m(self, argument)\u001b[0m\n\u001b[1;32m     80\u001b[0m output_structure \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39moutput_structure()\n\u001b[1;32m     81\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m jax\u001b[38;5;241m.\u001b[39mnamed_scope(\u001b[38;5;28mtype\u001b[39m(\u001b[38;5;28mself\u001b[39m)\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m):\n\u001b[0;32m---> 82\u001b[0m   result \u001b[38;5;241m=\u001b[39m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margument\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     83\u001b[0m shapecheck\u001b[38;5;241m.\u001b[39mcheck_structure(\n\u001b[1;32m     84\u001b[0m     result,\n\u001b[1;32m     85\u001b[0m     pattern\u001b[38;5;241m=\u001b[39moutput_structure,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     90\u001b[0m     ),\n\u001b[1;32m     91\u001b[0m )\n\u001b[1;32m     92\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m result\n",
      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/micrlhf-progress-a058ydGG-py3.12/lib/python3.12/site-packages/penzai/data_effects/side_input.py:146\u001b[0m, in \u001b[0;36mWithSideInputsFromInputTuple.__call__\u001b[0;34m(self, argument)\u001b[0m\n\u001b[1;32m    136\u001b[0m impls \u001b[38;5;241m=\u001b[39m {\n\u001b[1;32m    137\u001b[0m     tag: SideInputEffectImpl(_value\u001b[38;5;241m=\u001b[39mval, _handler_id\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mhandler_id)\n\u001b[1;32m    138\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m tag, val \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mzip\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mside_input_tags, side_inputs)\n\u001b[1;32m    139\u001b[0m }\n\u001b[1;32m    140\u001b[0m handled_body \u001b[38;5;241m=\u001b[39m (\n\u001b[1;32m    141\u001b[0m     selectors\u001b[38;5;241m.\u001b[39mselect(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mbody)\n\u001b[1;32m    142\u001b[0m     \u001b[38;5;241m.\u001b[39mat_instances_of(HandledSideInputRef)\n\u001b[1;32m    143\u001b[0m     \u001b[38;5;241m.\u001b[39mwhere(\u001b[38;5;28;01mlambda\u001b[39;00m ref: ref\u001b[38;5;241m.\u001b[39mhandler_id \u001b[38;5;241m==\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mhandler_id)\n\u001b[1;32m    144\u001b[0m     \u001b[38;5;241m.\u001b[39mapply(\u001b[38;5;28;01mlambda\u001b[39;00m ref: impls[ref\u001b[38;5;241m.\u001b[39mtag])\n\u001b[1;32m    145\u001b[0m )\n\u001b[0;32m--> 146\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mhandled_body\u001b[49m\u001b[43m(\u001b[49m\u001b[43minner_arg\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/micrlhf-progress-a058ydGG-py3.12/lib/python3.12/site-packages/penzai/nn/grouping.py:79\u001b[0m, in \u001b[0;36mSequential.__call__\u001b[0;34m(self, value)\u001b[0m\n\u001b[1;32m     77\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m i, layer \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28menumerate\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39msublayers):\n\u001b[1;32m     78\u001b[0m   \u001b[38;5;28;01mwith\u001b[39;00m jax\u001b[38;5;241m.\u001b[39mnamed_scope(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mi\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m):\n\u001b[0;32m---> 79\u001b[0m     value \u001b[38;5;241m=\u001b[39m \u001b[43mlayer\u001b[49m\u001b[43m(\u001b[49m\u001b[43mvalue\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     80\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m value\n",
      "File \u001b[0;32m~/micrlhf-progress/micrlhf/scan.py:35\u001b[0m, in \u001b[0;36mScanSequential.__call__\u001b[0;34m(self, inputs)\u001b[0m\n\u001b[1;32m     33\u001b[0m     h_ \u001b[38;5;241m=\u001b[39m pz\u001b[38;5;241m.\u001b[39mnx\u001b[38;5;241m.\u001b[39mwrap(h_\u001b[38;5;241m.\u001b[39munwrap(\u001b[38;5;241m*\u001b[39mh\u001b[38;5;241m.\u001b[39mnamed_axes\u001b[38;5;241m.\u001b[39mkeys()), \u001b[38;5;241m*\u001b[39mh\u001b[38;5;241m.\u001b[39mnamed_axes\u001b[38;5;241m.\u001b[39mkeys())\n\u001b[1;32m     34\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m h_, \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m---> 35\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mjax\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlax\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mscan\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m     36\u001b[0m \u001b[43m    \u001b[49m\u001b[43mscanner\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     37\u001b[0m \u001b[43m    \u001b[49m\u001b[43minputs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     38\u001b[0m \u001b[43m    \u001b[49m\u001b[43mlayer_nx\u001b[49m\u001b[43m)\u001b[49m[\u001b[38;5;241m0\u001b[39m]\n",
      "    \u001b[0;31m[... skipping hidden 9 frame]\u001b[0m\n",
      "File \u001b[0;32m~/micrlhf-progress/micrlhf/scan.py:32\u001b[0m, in \u001b[0;36mScanSequential.__call__.<locals>.scanner\u001b[0;34m(h, l)\u001b[0m\n\u001b[1;32m     30\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mscanner\u001b[39m(h, l):\n\u001b[1;32m     31\u001b[0m     l_ \u001b[38;5;241m=\u001b[39m eqx\u001b[38;5;241m.\u001b[39mcombine(l, layer_base, is_leaf\u001b[38;5;241m=\u001b[39mis_nx)\n\u001b[0;32m---> 32\u001b[0m     h_ \u001b[38;5;241m=\u001b[39m \u001b[43ml_\u001b[49m\u001b[43m(\u001b[49m\u001b[43mh\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     33\u001b[0m     h_ \u001b[38;5;241m=\u001b[39m pz\u001b[38;5;241m.\u001b[39mnx\u001b[38;5;241m.\u001b[39mwrap(h_\u001b[38;5;241m.\u001b[39munwrap(\u001b[38;5;241m*\u001b[39mh\u001b[38;5;241m.\u001b[39mnamed_axes\u001b[38;5;241m.\u001b[39mkeys()), \u001b[38;5;241m*\u001b[39mh\u001b[38;5;241m.\u001b[39mnamed_axes\u001b[38;5;241m.\u001b[39mkeys())\n\u001b[1;32m     34\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m h_, \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/micrlhf-progress-a058ydGG-py3.12/lib/python3.12/site-packages/penzai/nn/grouping.py:79\u001b[0m, in \u001b[0;36mSequential.__call__\u001b[0;34m(self, value)\u001b[0m\n\u001b[1;32m     77\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m i, layer \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28menumerate\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39msublayers):\n\u001b[1;32m     78\u001b[0m   \u001b[38;5;28;01mwith\u001b[39;00m jax\u001b[38;5;241m.\u001b[39mnamed_scope(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mi\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m):\n\u001b[0;32m---> 79\u001b[0m     value \u001b[38;5;241m=\u001b[39m \u001b[43mlayer\u001b[49m\u001b[43m(\u001b[49m\u001b[43mvalue\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     80\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m value\n",
      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/micrlhf-progress-a058ydGG-py3.12/lib/python3.12/site-packages/penzai/nn/combinators.py:59\u001b[0m, in \u001b[0;36mResidual.__call__\u001b[0;34m(self, value)\u001b[0m\n\u001b[1;32m     50\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__call__\u001b[39m(\u001b[38;5;28mself\u001b[39m, value: Any) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Any:\n\u001b[1;32m     51\u001b[0m \u001b[38;5;250m  \u001b[39m\u001b[38;5;124;03m\"\"\"Runs each of the sublayers in sequence, then adds back the original input.\u001b[39;00m\n\u001b[1;32m     52\u001b[0m \n\u001b[1;32m     53\u001b[0m \u001b[38;5;124;03m  Args:\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     57\u001b[0m \u001b[38;5;124;03m    The sum of the input to the residual block and the output of the child.\u001b[39;00m\n\u001b[1;32m     58\u001b[0m \u001b[38;5;124;03m  \"\"\"\u001b[39;00m\n\u001b[0;32m---> 59\u001b[0m   delta_value \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdelta\u001b[49m\u001b[43m(\u001b[49m\u001b[43mvalue\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     60\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m delta_value \u001b[38;5;241m+\u001b[39m value\n",
      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/micrlhf-progress-a058ydGG-py3.12/lib/python3.12/site-packages/penzai/nn/grouping.py:79\u001b[0m, in \u001b[0;36mSequential.__call__\u001b[0;34m(self, value)\u001b[0m\n\u001b[1;32m     77\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m i, layer \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28menumerate\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39msublayers):\n\u001b[1;32m     78\u001b[0m   \u001b[38;5;28;01mwith\u001b[39;00m jax\u001b[38;5;241m.\u001b[39mnamed_scope(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mi\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m):\n\u001b[0;32m---> 79\u001b[0m     value \u001b[38;5;241m=\u001b[39m \u001b[43mlayer\u001b[49m\u001b[43m(\u001b[49m\u001b[43mvalue\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     80\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m value\n",
      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/micrlhf-progress-a058ydGG-py3.12/lib/python3.12/site-packages/penzai/nn/attention.py:225\u001b[0m, in \u001b[0;36mKVCachingAttention.__call__\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m    222\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mkv_cache\u001b[38;5;241m.\u001b[39mset((new_key_cache, new_value_cache))\n\u001b[1;32m    224\u001b[0m \u001b[38;5;66;03m# Run the rest on the updated KV caches.\u001b[39;00m\n\u001b[0;32m--> 225\u001b[0m attn \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mquery_key_to_attn\u001b[49m\u001b[43m(\u001b[49m\u001b[43m(\u001b[49m\u001b[43mquery\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mnew_key_cache\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    226\u001b[0m output \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mattn_value_to_output((attn, new_value_cache))\n\u001b[1;32m    227\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m output\n",
      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/micrlhf-progress-a058ydGG-py3.12/lib/python3.12/site-packages/penzai/nn/grouping.py:79\u001b[0m, in \u001b[0;36mSequential.__call__\u001b[0;34m(self, value)\u001b[0m\n\u001b[1;32m     77\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m i, layer \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28menumerate\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39msublayers):\n\u001b[1;32m     78\u001b[0m   \u001b[38;5;28;01mwith\u001b[39;00m jax\u001b[38;5;241m.\u001b[39mnamed_scope(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mi\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m):\n\u001b[0;32m---> 79\u001b[0m     value \u001b[38;5;241m=\u001b[39m \u001b[43mlayer\u001b[49m\u001b[43m(\u001b[49m\u001b[43mvalue\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     80\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m value\n",
      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/micrlhf-progress-a058ydGG-py3.12/lib/python3.12/site-packages/penzai/core/layer.py:72\u001b[0m, in \u001b[0;36mchecked_layer_call.<locals>.wrapper\u001b[0;34m(self, argument)\u001b[0m\n\u001b[1;32m     70\u001b[0m \u001b[38;5;129m@functools\u001b[39m\u001b[38;5;241m.\u001b[39mwraps(func)\n\u001b[1;32m     71\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mwrapper\u001b[39m(\u001b[38;5;28mself\u001b[39m, argument: Any, \u001b[38;5;241m/\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Any:\n\u001b[0;32m---> 72\u001b[0m   input_vars \u001b[38;5;241m=\u001b[39m \u001b[43mshapecheck\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcheck_structure\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m     73\u001b[0m \u001b[43m      \u001b[49m\u001b[43margument\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     74\u001b[0m \u001b[43m      \u001b[49m\u001b[43mpattern\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43minput_structure\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     75\u001b[0m \u001b[43m      \u001b[49m\u001b[43merror_prefix\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m(\u001b[49m\n\u001b[1;32m     76\u001b[0m \u001b[43m          \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mError while checking the input of a layer of type\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\n\u001b[1;32m     77\u001b[0m \u001b[43m          \u001b[49m\u001b[38;5;124;43mf\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43m \u001b[39;49m\u001b[38;5;132;43;01m{\u001b[39;49;00m\u001b[38;5;28;43mtype\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;18;43m__qualname__\u001b[39;49m\u001b[38;5;132;43;01m}\u001b[39;49;00m\u001b[38;5;124;43m:\u001b[39;49m\u001b[38;5;130;43;01m\\n\u001b[39;49;00m\u001b[38;5;124;43m\"\u001b[39;49m\n\u001b[1;32m     78\u001b[0m \u001b[43m      \u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     79\u001b[0m \u001b[43m  \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     80\u001b[0m   output_structure \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39moutput_structure()\n\u001b[1;32m     81\u001b[0m   \u001b[38;5;28;01mwith\u001b[39;00m jax\u001b[38;5;241m.\u001b[39mnamed_scope(\u001b[38;5;28mtype\u001b[39m(\u001b[38;5;28mself\u001b[39m)\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m):\n",
      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/micrlhf-progress-a058ydGG-py3.12/lib/python3.12/site-packages/penzai/core/shapecheck.py:992\u001b[0m, in \u001b[0;36mcheck_structure\u001b[0;34m(value, pattern, known_vars, error_prefix)\u001b[0m\n\u001b[1;32m    987\u001b[0m   \u001b[38;5;28;01mif\u001b[39;00m unsolved_variable_constraints:\n\u001b[1;32m    988\u001b[0m     msgs\u001b[38;5;241m.\u001b[39mappend(\n\u001b[1;32m    989\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mThere may also be additional mismatches due to remaining\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    990\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m unsolved constraints.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    991\u001b[0m     )\n\u001b[0;32m--> 992\u001b[0m   \u001b[38;5;28;01mraise\u001b[39;00m StructureMismatchError(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;241m.\u001b[39mjoin(msgs))\n\u001b[1;32m    993\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(unsolved_variable_constraints) \u001b[38;5;241m<\u001b[39m \u001b[38;5;28mlen\u001b[39m(variable_constraints):\n\u001b[1;32m    994\u001b[0m   variable_constraints \u001b[38;5;241m=\u001b[39m unsolved_variable_constraints\n",
      "\u001b[0;31mStructureMismatchError\u001b[0m: Error while checking the input of a layer of type NamedEinsum:\nMismatch while checking structures:\nAt root[1]: Named shape mismatch between value {'batch': 16, 'projection': 128, 'layer': 80, 'kv_heads': 8, 'seq': 64} and pattern {'seq': var('I')['tkv'], 'kv_heads': var('I')['h'], 'projection': var('I')['p'], 'batch': var('B')['batch']:=16}:\n  Unexpected names in value's named shape: ['layer']\n  After inlining var('B') = {'batch': 16} from root[0]"
     ]
    }
   ],
   "source": [
    "from micrlhf.sampling import sample\n",
    "texts, cache = sample((llama, cache), tokenizer, prompt, batch_size=batch_size, max_seq_len=max_seq_len,\n",
    "       do_sample=True, return_only_completion=True)\n",
    "texts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div id=\"output_0957272383054c77a03350b2ea7bce9e\"><script> /* penzai.treescope rendering of a Python object (compressed) */ (()=>{ let observer; let lastStep = new Promise((resolve, reject) => { observer = new IntersectionObserver((entries) => { for (const entry of entries) { if (entry.isIntersecting) { resolve(); observer.disconnect(); return; } } }, {rootMargin: \"1000px\"}); }); window.treescope_decompress_enqueue = (encoded, destId) => { const previous = lastStep; const destElt = document.getElementById(destId); lastStep = (async () => { await previous; let blob = new Blob([ Uint8Array.from(atob(encoded), (m) => m.codePointAt(0)) ]); let reader = blob.stream().pipeThrough( new DecompressionStream(\"deflate\") ).pipeThrough( new TextDecoderStream(\"utf-8\") ).getReader(); let parts = []; while (true) { let step = await reader.read(); if (step.done) { break; } parts.push(step.value); } let newElt = document.createElement(\"div\"); newElt.innerHTML = parts.join(\"\"); destElt.parentNode.replaceChild(newElt, destElt); for (let oldScript of newElt.querySelectorAll(\"script\")) { let newScript = document.createElement(\"script\"); newScript.type = oldScript.type; newScript.textContent = oldScript.textContent; oldScript.parentNode.replaceChild(newScript, oldScript); } })(); requestAnimationFrame(() => { observer.observe(destElt); }); } })(); </script><div id=\"compress_html_4a23c81b279d4c82a35a00fb28e682d5\"><script>window.treescope_decompress_enqueue(\"eNrtXAlT20jT/isTpSrYCzaSL2wTU69sfJEA4UhIsu+WaySNDixLYjS2Mfvlv389I9+WOQJhl7xAVWyPevrunp7H7L4P2cgle2lGCQl1PyAd6vsM/Y0CP3SY43tlRImLmTMgu8j0PZYycc9xR2XU8z0/DLAO60PbYSQlPpRRQGHFdUKWEqxTbBTAqud7sKxhvWtRv+8ZKd13fVqOtu6i8SfNBQLg5xjMLiPTYUDmMeKxXdTD1HK8lEtMVkYZ3eYyPJKyiWPZsKKk87vox/vtyJz3oU6dgO0hx0SJoeMZ/nBmIapUKghUICYwMJJg6yoF+vvH7spyuhMQz3A8S9W5Z0Ig+/OvO8la2DNcztLru24MtUVYZ8n7FZSYOL3D/CSq7IGS23+gj77fhSBQxGyCBKXnGySN/thGLmFovHduqwgNCE9Ez95UkOHr/R44NK35xgi9e4fe8Cdp3cVh+BGCluYOx44XJqRFpaQk4q6ayOCbAkyBU90lPRGhHyCZ9aknHsLHGFtp3zvzfY8bOPRpd2wZiAwZaHYBS/zRwjJzdL4YEAqG97Cnk7TnDxPJqWkrT1Aq2vQeZTNc5bgcWA5l2iWexWxIDSTHZcTtMY3s5h4gbkhmqtt9j+t+p/DQdkzGTRI7+Jsf8HtfHRITSkqu+iRkquf0MKdoUNwjicivyWR8RJZVCfqhHcVm9z6OmyhRiVzxENfdX+1YxYkHu/pkn5iEUmKck14AaU/C1QQKXL44TXzYRkdnxCU686nqupDo473zTbDvdULIVCm5KwouwesLOCHfFAy5ofA6VzeU9PwBWaiaOR7rKiEx0ZXzN6l/Q3hxCM5R41tUNyFFjY2z5Dvg07xpOiVgxrgg52nhXZr3YSCOhIhP43VyzWqRsLnHs8VdoU5U60e83VAC3tJJDcrPSACHLU7A47Q2Vsy3LDdqIx3R/RkoJrocrBCXbSEyAO5jV/C0E5/TXTLiFSlRadJ6gHjO5xHfhDTl2emBgtJEjx8IzoPoIHhvOAMkNlaW+xpiWAOFyXVFkiXkeyAUtPfm6OL1TzDbCSeaS3DkRAepOMk6WNMoGYgcFAfb20Ixg2UZdBoT6H6PB2mOAosfrvUSCS57PkuUbUgvmoyhH5N7/Z5G6BxDuVRU8oUZQQhae9a8xKySV/ixmTbGVdQRkbV911jghHf47/j8F3aWkcOw6+h8swnkWAP/8MNozeQARi2QdSKXArXhhCB0NJkQlgnRHnKxRtxyWSNQiGROK1387MbKi8aElMLnhPE0Ac6fynI8MTxors/HjbUyhctXJRuYdkOCLcgPb3W3iNbSko3DxJ7guRfrhyjCuk30LkwkSfRHcqYD3xq/aUK/oKGo2TLa+G8mr+kb/6R6i5vWKll4BiV5HLngPg15AAPfAQ1ojFwnfDqxohSEoJRoEuG6HH8aqTPzePNelZJ2wo7p0JB1fK/D0z+mtG4rpXQmz6spNlTo0epHEV9WkVu1MPiLgv6R5v0aHGt0dH4IPVG2zPWicWNYEK1jV0/AxQPGSyW4FmfMkh5wQv4yVTSfQk8eq2L4DIRyNdBbIvNf7pSrPnY9GJ06cAMznWtgEvVr5wbatZwukh6nCm3MG/0QUw9Og84kUydN1TSxrmRjCAM4sSFBsMFHuRQd37p4OMZeGi+l5LTIk9lVryxuapimLIoNh08mSjZvEGsL+RR7FkEyaFfQbZgj4HMYTRrREorcvKrKSqY8meMXYj4xBrzBaF9nnWDEp4IOZoxO/bt4Hk5PXzmTyZt869AmHhzmrouDEKJ255H3YL1vkRAddEKGICLXAczfa2meQo84EWNDFwk6MJPMpsE4t6wnf5pwb6El3CO9OEmieyq87OVpKDqwD1Rgj972rPbeU/2FoN4eyPtIvZeL7uPL+wi7y6j/jKEunUzKfAJ7bRzCFYs62EVno57mw93yuM+4bANNblE1PxhtLE7LHkco3PHacIxZ5fiFIKR6GfWpmzAww2X+fHvom2ZmV8MhKeS2DLnUPLTUqip+2ieq6ot31dMh/NtqqGpdve2n2lNVq+t/MNr1am34TVXPv9UO1MN2taY2rOt266PNwuqhQ6xsY/9r5mO78G1wFvSdT4f5c+Xga/v0y+Hg4vCGfRo1GrXNC6t77lT3ZdvZP+kf1I3mpdzSts1B2wiuPhTsqwvHOekfek27ZX5m6udC9Yjm1Ebb69YL+ud+39s8zV/pYXc4MBvu9tW1VfeLlnYwbBaVlrrtqaf5j5QeKKeb1o18asjqgalYRzu1YfMyY8n+qH+6s9OrK4Vh62vp2LICct4d5Uhbu8nrGj1uMqxaJ+2j4T4OR+FJv93+elFvDNVPJ0H7m/F5e3vT2jnf+Zplsvnh05U6yAPPj+rRjno4VHvWzenZZv/7Gal/vc6YBf3mKHfaGuX7VfXDTfUyaARZp3VSq8vf+59yZzueWf1YbzUOe6qzWRzUM7an2Dub2pfh18thiw72m59r3qVZr1ts81j/7ro7+VLtYFgt2qXc4WHzLNv8rlq9dv6yelJi503SKtWr1XYzu2/lTre/6SNNbUJMv3zYVk+aWCWHNVdt3dSPre/MKlQ/WcfH7f1q1znJk0b1a63a0B05sKkfeJAbwff6vnKjdM/Mmsns0QevZeBG2DLlo16zflSoGurVly8BZuHZ955hYKeUMW9Kuc/O5VUh6NHCsf+tdubQZm9w0MyeXZxlG/WMXj0xzzdbrh80c41wmMfWVaHofCdnR25w4VVbbWIcUtK/uGrWespFg3bPzq7zmcLFRThUQaMkEugcS2yItN5IPsVBA9U90vqM+V5su5k9jrt1SEhaooJJGMaOeOL/EiVnSLs/2wSg2LWuAztFM+hBQ7JhhCkj7DHY7kCNG7OZYTLLrdxTFga5khjkFqyMLjgxVsTfV2c700McdnQYXMCx0/3YZAvX/8lF+zaZS3sWRc67Hg0wTaRSvNulsAeBFZhfcn6ZC+FgKJ8Gx9cRwRYpISLgsRTMZ36fPcyUqQYQGIcYbxY1ESLRG6cX+JRhb4W3Rv2uOBmD0ew2ebd357bN+XM6si9/XbECnNkCLBX7I9NQBUF383SUiCQsIuhgybnPEw/IxhrwtZAwoRR3KdcDD7HDkIcHjoWZT9PAOdB8TI30kDqMnMP1MTHjBZEY85oBb3ALSEhz/hYYI2HnTo9AYKag5sq+CUi6tPXHFsrIsixag44ZzP0JcfWLlzvnVGmm3CKMmYDwvkUN7LiQDMxHnPiNqEkY8D24KLkjGBlCRrDBc2Jz3ndj+HCGHfJUmoCHi0DR8i1E2nsftbH3jhf0GeIoa0USjUvzr6VYJuMeBw+j/rb3fltsXpRrAycKIrodBhcUwiShcapLRtDN7Yok7SVAYdgA28cqzG3fmOXkBvK9Gnd9ZeOWPBPgZnIDTYujIiVc3NMMLGDYsvg3KSGRvxVproqgOcZQ8k64XPNAKfzNDYD3UqTx3tSM6GXeCSvXbem259LDPPin/Jf0XLFeirFjVKQpBFsgOY3oWYz1opaTS6Rk6tmsntVk0zCLipZbUvKuiX2JfB6elvbeuWz3CPeIoVKKR0gzXR8zpZD4P7i+QxmWlcIWCqh/ScRXN2UlU9xCcNQSWi7KW6g76NhQQGEZVkNyxR8nUWJIcRBwoPkSX6cF3+Q7i+3eEtb1l6w4W9c9XII/pL2AeDfYgc5HSZo/AOdckzC9qMjM+sR6nWJ9OIbnpT0EneYJPLfkoKWX1Zq+v9seWjYzZ1UWmsgqDLMRG8+FlH/CMrl/Ic/F+1EpxOuKjL8ZXYzIMYfjiLHv6Czx97N33kUT79+Hl/b9G7tyXKlF31xBu3p7ndnZFQUWvV3UqowekSd/bgi+G3/FahB9uSbtKYVFkVu39GLoDHd2O+nZk2Vq55qseXfV99nu/blE9D+dSPdNq8enzawb/4LcmTG/K4EyxZefQfPWPjKNFli9mFwSJ/kvSCPB944MKsovP4HGdj4ydyZcXkzaTMa+X5A5E9Z3Jc/Lz52ZpY9MnzlGLyaD4LLwC5IHuD702LotJ7b+lXkjjHxkykQ8njRbopcfyV9SkVLct/3A5i06JQH/W0EDBSNm+x7ytUuEGZKvd8x8SZczesHMyTHSftP7zE9W8OpfBUh7XH4Hc0RhzK4SA/KU8jouaGa+aJBMTs4ZpbySL2aNDDblTCmfyUm//EYdhyjE/RlgBBJNsZwZ0sERDgFrcECjKN4nUTqdRjFoz0vqGbP4PSwP5/ZtPGlviEeGXsjR/SAfPn8B33lSKv8OUNgsmtAy5EJWLxo5Q8+UsmbBLGmY7CikJBuZV1D4FRR+BYVvL+TfHhRWfnKIUl5B4TV58huDwsqTgMLKKyi8Pnd+c1BYeTpQWHkFhden0W8LCitPAgorr6Dwusz5jUFh5alAYeUVFI5Pnt8TFFaeABRWXkFh73/hPvPcoHBWwTibz8vZQiGby+ZwkZRkPZshxYyp5w1TewWF/7Ge8VOgsPLsoPCLcOaDPPir3PaPOuFZ/nB68mI4g/X/gcG6/4tKIjn98/f/Bw/mJ9I=\", \"compress_html_4a23c81b279d4c82a35a00fb28e682d5\");</script><span style=\"color: #aaaaaa; font-family: monospace\">(Loading...)</span></div></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div id=\"compress_html_5b1b3ffb1acd4820b8d19c7015fcb02f\"><script>window.treescope_decompress_enqueue(\"eNrtXNtyGze2fc9XoNo1Fjkh22zexIvIqowzp2pe8jB1qs6DStWFBnaz22o22mhQF2fyHv9H8gH5hXyKv+QA3RQpShQtK4pcmVoqu0kBG5e118ICKMM6kekFK811RjNPpmWR8esJy1VOHkvlzIuVDiXFpDXJcDwQfBjFg5Gkbr/Tl+NBMBj1ZJfHne540O1785Oy4Ll9uv7mfqSMUcsszWmSK9PwY5VJHmUU5krSJOFlY57xiLL5bk1o1GKR1W0mIiFxTrLZZH9vsh9ZwaVM80XbqGLCOn4woOWUuRHaCaWLxExY4FdlP/lCZXbyQi2XlJvwMqE8pCs7PUny2aZTjTFhr3j1tRnV5srNaRsQZVycT1lknwutVrlsrysuuG6021GmxHld1LRRSkvSNcSguGKlylJ5KLLO8+eCf2J+cl2QtjM7DzUtlaEwURek2d1J35+mOs9E0hj1/+Zy3h+woDOue7zkZVgKrbLMNjfKj1NdmtCk4vzRvY77fuD67YxZ0Bv6x3s7rrJ/cLjmvQGN5nmZmlTlk3uDs6BkxEtqp3lbrUyLxSo37ctKQ3fqHKuJSEIrbUFhRrG5pUP3rRPiQCS7cbrqaRuoa3VuIkvDnZBCkaSZ1JRvJl9zyuZsN+TZJLtZ5JU6pmzJ9cICrYEInolGVySsbWfaF0nNxWfm+qdNbS3vemq1sKtZ7ZX3jjNUOGyus4T9nTUCC+d2k3W/pY2iZrO5ycG2afvxbW/GvVmEVctNlzel7XXxC2ZzayNVBtuay3RVrgEGtLToHgZWzfSPTeKzEv9xx+lqfp2FPcLpDgf/5Nse/5RtZ0vygZ3n5E29/1WbIRMZL8uZV+0/lQEZnRbrnfJm490BYLcTGY36Az7dR40dZurtdr2DwtbV4E7SvFgZZqzlz7wKTqSuvL1t1sht0zdVo93ut5uGsZom4zHJDW+f03XBTTLz7sxmx4TXPLBd//Qek5p15fsVz3K+pLDQFKdX3ryg/ANPfaOJSqEKsqoq7ExKu5vFPMuc0fuWgeocUr/8oMy/b3ovHOjGbv3Jm3XGbo175HqNVpbt/Iip/G1mt5jZ0XbQxJ4jbPZcVChcZcMkadk8qnPjimdeI+PLSHKmlbLqcM+mt2V8EzhhR3sij6ZsHcNzq0heb2NHb20D5vJu32+OW7tobmfv7hJku0vQOxB9L1TZLSzNeWbPDoWefZFE/J22+wQTlpadfOHNX7+66h5PX2dm+o5f+d9pza9ZFGeKm2DYCIYtFnRHLTbq2L/V+yb79PFne3hgv//mnqefPv5qX1u29Bf7esY+kFaTYHgcHh93w2DYeb0w03qM3bS17kvSTizjRUnSm7N9yd05UnovIJ/dND5eTHfaPV1ajxTaw9KxHymeKJxU7pVNvlpGboRxf9TpBcHgeGQ/kvT/ytSm8inE2lZfkdZqk3kisa7ty25oBz4Yegdl8GqvVu7tUM641oGVfx3YZg4JzUttWW5uuffBem/+fATUGeJRpOnCWrI146qC1ca5Mea1V1tHPSChr7iUHLKnLKaq3XMtp8cutTt5aiGXj0rbznFz+wMndueTuDdvfhkTdpoPnkk/t7ZOO2e+Kwm5WyVfw9xulu4f8b1be6RzgEfoceOPf87R7fAU77vQC9luBbRx6r7s5Os/vu/fvO2ctb5hm6/nCXEVz97pAyE2ZhsEiP8NEKuBQSog/gUgOozfbDGCR0CExYJUQITFgkdYLEgFxL+axdbCheMCIhwXpAIiDrXgERBhsYAIiwWPgAiLBamAWFts5bHfbD0WTAIiTBakAiLOseARFgtSAREWCx4BERYLUmGxuGIAiHBckAqIONSCR0CExQIiLBY8AiIsFqQCIq4YQKwwWZAKiDjHgkdAhMVCt7BY8AiIsFiQCoi4YgDlwnFBKiDiUAseAREWC1JhseAREGGxIBUQn+GKwfoHB7hxAIjwXJAKiDjWgkdAhMUCIiwWPAIiLBakAiJuHEC5cFyQCog41IJHQITFQrewWPAIiLBYkAqI+KUGECtMFqQCIs6x4BEQYbHQLSwWPAIiLBakAiKuGEC5cFyQCog41IJHQITFglRYLHgERFgsSAVEXDEAkzBZkAqIOMeCR0CExYJUWCx4BERYLEgFRFwxAK2ACMcFRBxqwSMgwmJBKiDCYgERFgtSAfFFrhictZg01wXNojhT3ATD5smbsuD5/EteZHoxP7EPVprrjGaeTMsi49cTlqucPJbKmRcrHUqKSWuS4ZD6EYke52IU9TtjGsei1xO9qBPLeBREfW9+UvfsnkxkvCxn3mVCeShUlvGitH3wXIa5MqFWq1wanRbebriNtEPyKNJ04c1fZ2b6A1+S/E5rfs1uwDb+wyJuRDIJhi1WaPWOhElVPgm6oxazCEhPRjZT5xdhQlyWE1ta0ntX3WSfPv7c8Tvs99/c8/TTx1/ta8uW/mJfz9gH0sr2ehweH3fDYNhhjUvNiyLNF+wdv/KraTRfL8z0Ti7vIaYrWyItYIvmIaxV4EOV71c8yy30sNAUp1fevKD8A099oTT5rsLm8opKf3ci22Q1Hp7T3pQLtVxSbrw5Y6++QqKfV6m9gPPeYNDpDYe9fq/PRzTuiF6XRt1YDGQcPazUe3SsR2y3o0yJ83aVrAl7JaNRf8CnN8WR0pJ0uxQ8owkLpndSHKtM8igjK31Jti7jEWXzkzQvVoZV69gTCYnzSF15e9uERi0WmWv6pmq0231ie9BZmp+HhusFGY9Jbnj7nK4LbpKZt49wmyXbgphRRfWaiCS0QYLCjGLzHEo1mqgUqiCr2cLOpDQUxjzLIi7O78pWmX/f9F440I27Clhn7Na4R67XaGWMyo+Yyt9mqTifHW0HTazWbfZcVChcZcMkadk8qnPjimdeI+PLSHKmlTKT6tn0toxvAifsaE/k0ZStY3huLY1XK4MdvbUNmMu7fb/R2cNuURqbDkuGSNJMarLla3ZqSXkHou+FKp0u0pxnoaZCz75IIv5O270OUVp28oU15VdX3eOps+aNJW4NwxlF5Q7OF0bV+y+1Amev9Ri7aWsd2FesbX3Wir0XkM9uGh8vpjvtni6tRwrtYelYL32icNL9G0u+WkZuhHF/1OkFweB41O11O39lalP5FGJtq69Ia7XJPJFY1/ZlN7SdI0l46PByp+7VXq3c26Gcca0DK/86sM0cPDqltiw3t9z7YL03fz4C7p+TqwpWG+fGmNde/QUH1pddSg7ZUxZT1e65ltNjl9qdPLWQy0elbee4GSmLf3nnxGlnktgjZ/PLmDj0MfJza+s0OPNdScjdKvka5nazdP+I793aI50DPEKPG3/8c45uh6d434VeyHbrD8Sn+L8YgIifRoNUQMQ/+IFHWCwsFhBhseAREGGxIBUWi/+LAYhwXJAKiDjUgkdAhMUCIiwWPAIiLBakAiJ+3SPECpMFqYCIcyx4BERYLHQLiwWPgAiLBamAiCsGUC4cF6QCIg614BEQYbEgFRYLHgERFgtSARFXDCBWmCxIBUScY8EjIMJiQSosFjwCIiwWpAIirhiAVkCE4wIiDrXgERBhsSAVEGGxgAiLBamA+GJXDNY/OMCNA0CE54JUQMSxFjwCIiwWEGGx4BEQYbEgFRBx4wDKheOCVEDEoRY8AiIsFrqFxYJHQITFglRAxC81gFhhsiAVEHGOBY+ACIuFbmGx4BEQYbEgFRBxxQDKheOCVEDEoRY8AiIsFqTCYsEjIMJiQSog4ooBmITJglRAxDkWPAIiLBakwmLBIyDCYkEqIOKKAWgFRDguIOJQCx4BERYLUgERFguIsFiQCogvcsXgrMWkuS5oFsWZ4iYYNk/elAXP51/yItOL+Yl9sNJcZzTzZFoWGb+esFzl5LFUzrxY6VBSTFqTDONRPB6IzrAnRrIvRXfci4fxOOJ0HNC4I7ve/KTu2T2ZyHhZzrzLhPJQqCzjRWn74LkMc2VCrVa5NDotvN1wG2mH5FGk6cKbv87M9Ae+JPmd1vya3YBt/IdF3IhkEgxbrNDqHQmTqnwSdEctZhGQnoxsps4vwoS4LCe2tKT3rrrJPn38ueN32O+/uefpp4+/2teWLf3Fvp6xD6SV7fU4PD7uhsGwwxqXmhdFmi/YO37lV9Novl6Y6Z1c3kNMV7ZEWsAWzUNYq8CHKt+veJZb6GGhKU6vvHlB+Qee+kJp8l2FzeUVlf7uRLbJajw8p70pF2q5pNx4c8ZefYVEP0KppbBZMvNGo8lmc/ajUHlpmOF6QeZfsmQzduptpGp1yodRPBhJ6vY7fTkeBINRT3Z53OmOB92+12Lb4CH1IxI9zsUo6nfGNI5Fryd6USeW8SiIdoN7Aee9waDTGw57/V6fj6z0Ra9Lo24sBjKOdoI/u2LOpjUMrZSxCC7TXKpL32iiUqiCfIst3HwXuqgGk0qsHFWu8p8Zubf/uP6XbGxScdo5azanrO451nzhQmzvm4ZCEzf0/frb/1lHNDZt3DhvM2sCttFNe98uA8rl2yTNZMMF+MJF/KAkNYxekRvRugVrZGRYaht2pvblZEuQn1G+MIkt/fbbJvuRubi68vYoD2JKz+wArk2pVlrQ99a2Zg+movItj33L7rSvv/U1WZ8T9H+pSRqb7vw41eVNLxVM22Bbq2mpLsjl6KcqPzudbBLmqpsuyEq3luuNbO9RS/n7Fa3o+7VW/peWtj9DZaO5afv/UezcuQ==\", \"compress_html_5b1b3ffb1acd4820b8d19c7015fcb02f\");</script></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div id=\"output_dest_0957272383054c77a03350b2ea7bce9e\"><script> (()=>{ const output = document.getElementById(\"output_0957272383054c77a03350b2ea7bce9e\"); const dest = document.getElementById(\"output_dest_0957272383054c77a03350b2ea7bce9e\"); dest.parentNode.replaceChild(output, dest); })(); </script></div>"
      ],
      "text/plain": [
       "(<NamedArray bfloat16(| batch:16, projection:128, layer:80, kv_heads:8, seq:128) ≈0.0 ±0.0 [≥0.0, ≤0.0] zero:167_772_160 (wrapping jax.Array)>,\n",
       " <NamedArray bfloat16(| batch:16, projection:128, layer:80, kv_heads:8, seq:128) ≈0.0 ±0.0 [≥0.0, ≤0.0] zero:167_772_160 (wrapping jax.Array)>)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "next(iter(cache.kv_caches.values()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "micrlhf-progress-a058ydGG-py3.12",
   "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}