examples/llama2_inference/llama2.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c0c9f034-7af1-4dc2-bbfb-5bb9e27c07ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
    "import torch\n",
    "from torch.utils import _pytree as pytree\n",
    "from shark_turbine.aot import *\n",
    "from iree.compiler.ir import Context\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4d92bb47-2b93-4f32-a445-c0ad2adc37ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "#set some config values\n",
    "\n",
    "hf_auth_token = \"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\"\n",
    "hf_model_name = \"meta-llama/Llama-2-7b-chat-hf\"\n",
    "state_schema_path = \"llama2_state_schema.json\"\n",
    "with open(state_schema_path, \"r+\") as f:\n",
    "    state_schema = pytree.treespec_loads(f.read())\n",
    "prompt = \"\"\"\n",
    "<s>[INST] <<SYS>>\n",
    "Be concise. You are a helpful, respectful and honest assistant. If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information. <</SYS>> hi what are you? [/INST]\n",
    "\"\"\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d4664585-5e15-45c7-8c5c-c8eaf6381435",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/transformers/models/auto/tokenization_auto.py:640: FutureWarning: The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers.\n",
      "  warnings.warn(\n",
      "/home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/transformers/models/auto/auto_factory.py:479: FutureWarning: The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5e411acda19c4228b008ff622bdf110e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/transformers/utils/hub.py:374: FutureWarning: The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "#Initialize the tokenizer and base model from Huggingface\n",
    "tokenizer = AutoTokenizer.from_pretrained(\n",
    "    hf_model_name,\n",
    "    use_fast=False,\n",
    "    use_auth_token=hf_auth_token,\n",
    ")\n",
    "mod = AutoModelForCausalLM.from_pretrained(\n",
    "    hf_model_name,\n",
    "    torch_dtype=torch.float,\n",
    "    use_auth_token=hf_auth_token,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "48cdee7a-0a83-4299-ae6e-162e10f9ecdc",
   "metadata": {},
   "outputs": [],
   "source": [
    "#get some sample input values\n",
    "initial_input = tokenizer(prompt, return_tensors=\"pt\")\n",
    "example_input_id = initial_input.input_ids\n",
    "\n",
    "\n",
    "#Define some dim sizes to make the model code readable\n",
    "BATCH_SIZE = 1 #Note: only 1 supported currently.\n",
    "HEADS = 32\n",
    "HIDDEN_DIM = 128\n",
    "MAX_STEP_SEQ = 4095\n",
    "#This is a maximum size past key value tensor, representing the full context window of the model. \n",
    "#We do this to make it easier for the compiler to reason about memory.\n",
    "global_pkv = torch.zeros(\n",
    "    size=(HEADS * 2, BATCH_SIZE, HEADS, MAX_STEP_SEQ, HIDDEN_DIM),\n",
    "    dtype=torch.float32,\n",
    ")\n",
    "seq_step = AbstractIndex"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a220e7c8-47d6-4898-b876-24988decb328",
   "metadata": {},
   "outputs": [],
   "source": [
    "#define some helper functions for manipulating the pkv state\n",
    "\n",
    "def slice_up_to_step(global_pkv, seq_step, heads, hidden_dim):\n",
    "    \"\"\"\n",
    "    global_pkv: the global pkv tensor\n",
    "    seq_step: the current token index of the model\n",
    "    heads: the number of attn heads\n",
    "    hidden_dim: feature dimension size\n",
    "    takes the global_pkv tensor and gets the seq_step pair for each head\n",
    "    \"\"\"\n",
    "    all_pkv_tensors = []\n",
    "    for i in range(heads * 2):\n",
    "        sliced = IREE.tensor_slice(\n",
    "            global_pkv, i, 0, (0, heads), (0, seq_step), (0, hidden_dim)\n",
    "        )  # sequence context dim\n",
    "        all_pkv_tensors.append(\n",
    "            IREE.tensor_reshape(sliced, 1, heads, seq_step, hidden_dim)\n",
    "        )\n",
    "\n",
    "    return all_pkv_tensors\n",
    "\n",
    "\n",
    "def update_state(global_pkv, state_updates, seq_step, heads, hidden_dim):\n",
    "    \"\"\"\n",
    "    global_pkv: the global pkv tensor\n",
    "    state_updates: the state updates output by a forward pass of the model\n",
    "    seq_step: the current token index of the model\n",
    "    heads: the number of attn heads\n",
    "    hidden_dim: feature dimension size\n",
    "    updates the global state of the model at seq_step with state_updates\n",
    "    \"\"\"\n",
    "    all_updates = []\n",
    "    for i in range(heads * 2):\n",
    "        #expand dim in state updates to match the rank of global_pkv\n",
    "        update = IREE.tensor_reshape(\n",
    "            state_updates[i], 1, 1, heads, 1, hidden_dim\n",
    "        )\n",
    "        all_updates.append(\n",
    "            IREE.tensor_update(global_pkv, update, i, 0, 0, seq_step, 0)\n",
    "        )\n",
    "    return all_updates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e3cfb777-8f95-40f0-9dc7-4bf81cf91b40",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Set up our CompiledModule class\n",
    "class StateUpdateModule(CompiledModule):\n",
    "    #export_parameters makes the params of the model global, so that multiple \n",
    "    # exported functions can access them without duplicating the constants.\n",
    "    # We set external parameters mode to increase readability of the IR.\n",
    "    params = export_parameters(mod, external=True)\n",
    "    # export our global pkv tensor, making sure to set mutable=True so it \n",
    "    # can be modified by our exported functions\n",
    "    global_state = export_global(abstractify(global_pkv), mutable=True)\n",
    "    # this is our sequences step, its the current token index for the model\n",
    "    global_seq_step = export_global(AbstractIndex, mutable=True)\n",
    "\n",
    "    #set up our stateless jittable functions\n",
    "    @jittable\n",
    "    def initialize(input_ids):\n",
    "        result = mod.forward(input_ids)\n",
    "        state1_flat, _ = pytree.tree_flatten(result.past_key_values)\n",
    "        token = torch.argmax(result.logits[:, -1, :], dim=1)\n",
    "        token = token[None, :]\n",
    "        return token, *state1_flat\n",
    "\n",
    "    @jittable\n",
    "    def forward(token0: torch.Tensor, *state0_flat):\n",
    "        # Unpad the states.\n",
    "        state0 = pytree.tree_unflatten(state0_flat, state_schema)\n",
    "        result = mod.forward(token0, past_key_values=state0)\n",
    "        state1_flat, _ = pytree.tree_flatten(result.past_key_values)\n",
    "        # extract only the newest pkvs for each head\n",
    "        state1_flat = [x[:, :, -1:, :] for x in state1_flat]\n",
    "        token = torch.argmax(result.logits[:, -1, :], dim=1)\n",
    "        token = token[None, :]\n",
    "        return token, *state1_flat\n",
    "    \n",
    "\n",
    "    \"\"\"\n",
    "    run_initialize is doing the \"first forward\" pass of the llama model, \n",
    "    which computes the initial past_key_values for the input tokens (prompt)\n",
    "    \"\"\"\n",
    "    def run_initialize(\n",
    "        self, x=AbstractTensor(BATCH_SIZE, None, dtype=torch.int64)\n",
    "    ):\n",
    "        \"\"\" \n",
    "        We need to set constraints for the dynamic dimensions of the input.\n",
    "        This may look like magic, but if you try and export without it, \n",
    "        torch compile will actually try to calculate the constraints for you and tell you what to add\n",
    "        \"\"\"\n",
    "        init_const = [x.dynamic_dim(1) < MAX_STEP_SEQ]\n",
    "        token, *state = self.initialize(x, constraints=init_const)\n",
    "        updates = []\n",
    "        #initialize the token index to the 3rd dim of the pkv output\n",
    "        self.global_seq_step = IREE.tensor_dim(\n",
    "            state[0], 3\n",
    "        )\n",
    "        #iterate through the past key values and update global state\n",
    "        for i in range(HEADS * 2):\n",
    "            slice_of_state = IREE.tensor_reshape(\n",
    "                state[i], 1, 1, HEADS, self.global_seq_step, HIDDEN_DIM\n",
    "            )\n",
    "            updates.append(\n",
    "                IREE.tensor_update(\n",
    "                    self.global_state, slice_of_state, i, 0, 0, 0, 0\n",
    "                )\n",
    "            )\n",
    "        return token\n",
    "\n",
    "    def run_forward(self, x=AbstractTensor(1, None, dtype=torch.int64)):\n",
    "        state_arg = slice_up_to_step(\n",
    "            self.global_state, self.global_seq_step, HEADS, HIDDEN_DIM\n",
    "        )\n",
    "        forw_const = [state_arg[0].dynamic_dim(2) < MAX_STEP_SEQ] + [\n",
    "            x.dynamic_dim(2) == (state_arg[0].dynamic_dim(2))\n",
    "            for x in state_arg[1:]\n",
    "        ]\n",
    "        token, *state_update = self.forward(\n",
    "            x, *state_arg, constraints=forw_const\n",
    "        )\n",
    "        self.global_seq_step = self.global_seq_step + 1\n",
    "        res = update_state(\n",
    "            self.global_state,\n",
    "            state_update,\n",
    "            self.global_seq_step,\n",
    "            HEADS,\n",
    "            HIDDEN_DIM,\n",
    "        )\n",
    "        return token"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "10f78369-765f-439c-9772-93449dce4ba6",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2023-10-09 18:48:30,657] torch.fx.experimental.symbolic_shapes: [INFO] create_env\n",
      "[2023-10-09 18:48:30,660] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s0 = 2 for input0.size()[1]\n",
      "/home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/overrides.py:110: UserWarning: 'has_cuda' is deprecated, please use 'torch.backends.cuda.is_built()'\n",
      "  torch.has_cuda,\n",
      "/home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/overrides.py:111: UserWarning: 'has_cudnn' is deprecated, please use 'torch.backends.cudnn.is_available()'\n",
      "  torch.has_cudnn,\n",
      "/home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/overrides.py:117: UserWarning: 'has_mps' is deprecated, please use 'torch.backends.mps.is_built()'\n",
      "  torch.has_mps,\n",
      "/home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/overrides.py:118: UserWarning: 'has_mkldnn' is deprecated, please use 'torch.backends.mkldnn.is_available()'\n",
      "  torch.has_mkldnn,\n",
      "[2023-10-09 18:48:30,910] torch.fx.experimental.symbolic_shapes: [INFO] eval s0 <= 4096 [guard added] (transformers/models/llama/modeling_llama.py:119 in forward)\n",
      "[2023-10-09 18:48:30,918] torch.fx.experimental.symbolic_shapes: [INFO] eval Ne(s0, 4096) [guard added] (transformers/models/llama/modeling_llama.py:123 in forward)\n",
      "[2023-10-09 18:48:43,109] [0/0] torch.fx.experimental.symbolic_shapes: [INFO] create_env\n",
      "[2023-10-09 18:48:43,124] [0/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s0 = 2 for L['arg0_1'].size()[1]\n",
      "[2023-10-09 18:49:00,136] [0/0] torch.fx.experimental.symbolic_shapes: [INFO] produce_guards\n",
      "[2023-10-09 18:49:15,183] torch.fx.experimental.symbolic_shapes: [INFO] create_env\n",
      "[2023-10-09 18:49:15,186] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s0 = 2 for input0.size()[1]\n",
      "[2023-10-09 18:49:15,189] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s1 = 32 for input1.size()[1]\n",
      "[2023-10-09 18:49:15,191] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s2 = 128 for input1.size()[3]\n",
      "[2023-10-09 18:49:16,275] torch.fx.experimental.symbolic_shapes: [INFO] eval 2*s0 < 9223372036854775807 [guard added] (_decomp/decompositions.py:725 in slice_forward)\n",
      "[2023-10-09 18:49:16,412] torch.fx.experimental.symbolic_shapes: [INFO] eval 2*s0 <= 4096 [guard added] (transformers/models/llama/modeling_llama.py:119 in forward)\n",
      "[2023-10-09 18:49:16,423] torch.fx.experimental.symbolic_shapes: [INFO] eval Ne(2*s0, 4096) [guard added] (transformers/models/llama/modeling_llama.py:123 in forward)\n",
      "[2023-10-09 18:49:16,540] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s1, 32) [guard added] (_prims/__init__.py:1873 in _cat_meta)\n",
      "[2023-10-09 18:49:16,545] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s2, 128) [guard added] (_prims/__init__.py:1873 in _cat_meta)\n",
      "[2023-10-09 18:49:32,246] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_env\n",
      "[2023-10-09 18:49:32,263] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s0 = 2 for L['arg0_1'].size()[1]\n",
      "[2023-10-09 18:49:32,276] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s1 = 2 for L['arg1_1'].size()[2]\n",
      "[2023-10-09 18:49:32,286] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s2 = 2 for L['arg2_1'].size()[2]\n",
      "[2023-10-09 18:49:32,291] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s3 = 2 for L['arg3_1'].size()[2]\n",
      "[2023-10-09 18:49:32,301] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s4 = 2 for L['arg4_1'].size()[2]\n",
      "[2023-10-09 18:49:32,310] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s5 = 2 for L['arg5_1'].size()[2]\n",
      "[2023-10-09 18:49:32,319] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s6 = 2 for L['arg6_1'].size()[2]\n",
      "[2023-10-09 18:49:32,328] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s7 = 2 for L['arg7_1'].size()[2]\n",
      "[2023-10-09 18:49:32,338] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s8 = 2 for L['arg8_1'].size()[2]\n",
      "[2023-10-09 18:49:32,348] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s9 = 2 for L['arg9_1'].size()[2]\n",
      "[2023-10-09 18:49:32,357] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s10 = 2 for L['arg10_1'].size()[2]\n",
      "[2023-10-09 18:49:32,366] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s11 = 2 for L['arg11_1'].size()[2]\n",
      "[2023-10-09 18:49:32,375] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s12 = 2 for L['arg12_1'].size()[2]\n",
      "[2023-10-09 18:49:32,384] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s13 = 2 for L['arg13_1'].size()[2]\n",
      "[2023-10-09 18:49:32,393] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s14 = 2 for L['arg14_1'].size()[2]\n",
      "[2023-10-09 18:49:32,401] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s15 = 2 for L['arg15_1'].size()[2]\n",
      "[2023-10-09 18:49:32,410] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s16 = 2 for L['arg16_1'].size()[2]\n",
      "[2023-10-09 18:49:32,420] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s17 = 2 for L['arg17_1'].size()[2]\n",
      "[2023-10-09 18:49:32,430] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s18 = 2 for L['arg18_1'].size()[2]\n",
      "[2023-10-09 18:49:32,439] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s19 = 2 for L['arg19_1'].size()[2]\n",
      "[2023-10-09 18:49:32,448] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s20 = 2 for L['arg20_1'].size()[2]\n",
      "[2023-10-09 18:49:32,457] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s21 = 2 for L['arg21_1'].size()[2]\n",
      "[2023-10-09 18:49:32,468] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s22 = 2 for L['arg22_1'].size()[2]\n",
      "[2023-10-09 18:49:32,477] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s23 = 2 for L['arg23_1'].size()[2]\n",
      "[2023-10-09 18:49:32,486] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s24 = 2 for L['arg24_1'].size()[2]\n",
      "[2023-10-09 18:49:32,495] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s25 = 2 for L['arg25_1'].size()[2]\n",
      "[2023-10-09 18:49:32,505] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s26 = 2 for L['arg26_1'].size()[2]\n",
      "[2023-10-09 18:49:32,513] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s27 = 2 for L['arg27_1'].size()[2]\n",
      "[2023-10-09 18:49:32,523] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s28 = 2 for L['arg28_1'].size()[2]\n",
      "[2023-10-09 18:49:32,532] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s29 = 2 for L['arg29_1'].size()[2]\n",
      "[2023-10-09 18:49:32,541] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s30 = 2 for L['arg30_1'].size()[2]\n",
      "[2023-10-09 18:49:32,550] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s31 = 2 for L['arg31_1'].size()[2]\n",
      "[2023-10-09 18:49:32,559] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s32 = 2 for L['arg32_1'].size()[2]\n",
      "[2023-10-09 18:49:32,569] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s33 = 2 for L['arg33_1'].size()[2]\n",
      "[2023-10-09 18:49:32,578] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s34 = 2 for L['arg34_1'].size()[2]\n",
      "[2023-10-09 18:49:32,588] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s35 = 2 for L['arg35_1'].size()[2]\n",
      "[2023-10-09 18:49:32,597] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s36 = 2 for L['arg36_1'].size()[2]\n",
      "[2023-10-09 18:49:32,606] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s37 = 2 for L['arg37_1'].size()[2]\n",
      "[2023-10-09 18:49:32,616] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s38 = 2 for L['arg38_1'].size()[2]\n",
      "[2023-10-09 18:49:32,626] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s39 = 2 for L['arg39_1'].size()[2]\n",
      "[2023-10-09 18:49:32,636] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s40 = 2 for L['arg40_1'].size()[2]\n",
      "[2023-10-09 18:49:32,645] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s41 = 2 for L['arg41_1'].size()[2]\n",
      "[2023-10-09 18:49:32,655] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s42 = 2 for L['arg42_1'].size()[2]\n",
      "[2023-10-09 18:49:32,664] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s43 = 2 for L['arg43_1'].size()[2]\n",
      "[2023-10-09 18:49:32,673] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s44 = 2 for L['arg44_1'].size()[2]\n",
      "[2023-10-09 18:49:32,682] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s45 = 2 for L['arg45_1'].size()[2]\n",
      "[2023-10-09 18:49:32,693] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s46 = 2 for L['arg46_1'].size()[2]\n",
      "[2023-10-09 18:49:32,702] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s47 = 2 for L['arg47_1'].size()[2]\n",
      "[2023-10-09 18:49:32,710] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s48 = 2 for L['arg48_1'].size()[2]\n",
      "[2023-10-09 18:49:32,719] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s49 = 2 for L['arg49_1'].size()[2]\n",
      "[2023-10-09 18:49:32,728] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s50 = 2 for L['arg50_1'].size()[2]\n",
      "[2023-10-09 18:49:32,737] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s51 = 2 for L['arg51_1'].size()[2]\n",
      "[2023-10-09 18:49:32,747] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s52 = 2 for L['arg52_1'].size()[2]\n",
      "[2023-10-09 18:49:32,756] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s53 = 2 for L['arg53_1'].size()[2]\n",
      "[2023-10-09 18:49:32,765] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s54 = 2 for L['arg54_1'].size()[2]\n",
      "[2023-10-09 18:49:32,775] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s55 = 2 for L['arg55_1'].size()[2]\n",
      "[2023-10-09 18:49:32,784] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s56 = 2 for L['arg56_1'].size()[2]\n",
      "[2023-10-09 18:49:32,793] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s57 = 2 for L['arg57_1'].size()[2]\n",
      "[2023-10-09 18:49:32,802] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s58 = 2 for L['arg58_1'].size()[2]\n",
      "[2023-10-09 18:49:32,812] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s59 = 2 for L['arg59_1'].size()[2]\n",
      "[2023-10-09 18:49:32,821] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s60 = 2 for L['arg60_1'].size()[2]\n",
      "[2023-10-09 18:49:32,830] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s61 = 2 for L['arg61_1'].size()[2]\n",
      "[2023-10-09 18:49:32,839] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s62 = 2 for L['arg62_1'].size()[2]\n",
      "[2023-10-09 18:49:32,849] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s63 = 2 for L['arg63_1'].size()[2]\n",
      "[2023-10-09 18:49:32,858] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s64 = 2 for L['arg64_1'].size()[2]\n",
      "[2023-10-09 18:49:33,025] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval s0 + s1 < 9223372036854775807 [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:26 in forward (_decomp/decompositions.py:725 in slice_forward)\n",
      "[2023-10-09 18:49:33,234] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval s0 + s1 <= 4096 [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:72 in forward (_decomp/decompositions.py:725 in slice_forward)\n",
      "[2023-10-09 18:49:33,409] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s2, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:118 in forward (_meta_registrations.py:3515 in common_meta_baddbmm_bmm)\n",
      "[2023-10-09 18:49:33,707] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval s0 + s3 <= 4096 [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:189 in forward (_decomp/decompositions.py:725 in slice_forward)\n",
      "[2023-10-09 18:49:33,845] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s3, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:228 in forward (_subclasses/fake_tensor.py:740 in infer_size)\n",
      "[2023-10-09 18:49:33,878] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s4, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:235 in forward (_meta_registrations.py:3515 in common_meta_baddbmm_bmm)\n",
      "[2023-10-09 18:49:34,188] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval s0 + s5 <= 4096 [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:306 in forward (_decomp/decompositions.py:725 in slice_forward)\n",
      "[2023-10-09 18:49:34,326] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s5, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:345 in forward (_subclasses/fake_tensor.py:740 in infer_size)\n",
      "[2023-10-09 18:49:34,359] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s6, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:352 in forward (_meta_registrations.py:3515 in common_meta_baddbmm_bmm)\n",
      "[2023-10-09 18:49:34,661] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval s0 + s7 <= 4096 [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:423 in forward (_decomp/decompositions.py:725 in slice_forward)\n",
      "[2023-10-09 18:49:34,800] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s7, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:462 in forward (_subclasses/fake_tensor.py:740 in infer_size)\n",
      "[2023-10-09 18:49:34,832] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s8, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:469 in forward (_meta_registrations.py:3515 in common_meta_baddbmm_bmm)\n",
      "[2023-10-09 18:49:35,130] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval s0 + s9 <= 4096 [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:540 in forward (_decomp/decompositions.py:725 in slice_forward)\n",
      "[2023-10-09 18:49:35,271] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s9, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:579 in forward (_subclasses/fake_tensor.py:740 in infer_size)\n",
      "[2023-10-09 18:49:35,305] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s10, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:586 in forward (_meta_registrations.py:3515 in common_meta_baddbmm_bmm)\n",
      "[2023-10-09 18:49:35,611] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval s0 + s11 <= 4096 [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:657 in forward (_decomp/decompositions.py:725 in slice_forward)\n",
      "[2023-10-09 18:49:35,762] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s11, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:696 in forward (_subclasses/fake_tensor.py:740 in infer_size)\n",
      "[2023-10-09 18:49:35,795] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s12, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:703 in forward (_meta_registrations.py:3515 in common_meta_baddbmm_bmm)\n",
      "[2023-10-09 18:49:36,107] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval s0 + s13 <= 4096 [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:774 in forward (_decomp/decompositions.py:725 in slice_forward)\n",
      "[2023-10-09 18:49:36,249] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s13, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:813 in forward (_subclasses/fake_tensor.py:740 in infer_size)\n",
      "[2023-10-09 18:49:36,282] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s14, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:820 in forward (_meta_registrations.py:3515 in common_meta_baddbmm_bmm)\n",
      "[2023-10-09 18:49:36,589] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval s0 + s15 <= 4096 [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:891 in forward (_decomp/decompositions.py:725 in slice_forward)\n",
      "[2023-10-09 18:49:36,734] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s15, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:930 in forward (_subclasses/fake_tensor.py:740 in infer_size)\n",
      "[2023-10-09 18:49:36,768] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s16, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:937 in forward (_meta_registrations.py:3515 in common_meta_baddbmm_bmm)\n",
      "[2023-10-09 18:49:37,105] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval s0 + s17 <= 4096 [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:1008 in forward (_decomp/decompositions.py:725 in slice_forward)\n",
      "[2023-10-09 18:49:37,249] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s17, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:1047 in forward (_subclasses/fake_tensor.py:740 in infer_size)\n",
      "[2023-10-09 18:49:37,286] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s18, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:1054 in forward (_meta_registrations.py:3515 in common_meta_baddbmm_bmm)\n",
      "[2023-10-09 18:49:37,595] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval s0 + s19 <= 4096 [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:1125 in forward (_decomp/decompositions.py:725 in slice_forward)\n",
      "[2023-10-09 18:49:37,744] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s19, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:1164 in forward (_subclasses/fake_tensor.py:740 in infer_size)\n",
      "[2023-10-09 18:49:37,778] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s20, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:1171 in forward (_meta_registrations.py:3515 in common_meta_baddbmm_bmm)\n",
      "[2023-10-09 18:49:38,090] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval s0 + s21 <= 4096 [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:1242 in forward (_decomp/decompositions.py:725 in slice_forward)\n",
      "[2023-10-09 18:49:38,238] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s21, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:1281 in forward (_subclasses/fake_tensor.py:740 in infer_size)\n",
      "[2023-10-09 18:49:38,272] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s22, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:1288 in forward (_meta_registrations.py:3515 in common_meta_baddbmm_bmm)\n",
      "[2023-10-09 18:49:38,584] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval s0 + s23 <= 4096 [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:1359 in forward (_decomp/decompositions.py:725 in slice_forward)\n",
      "[2023-10-09 18:49:38,734] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s23, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:1398 in forward (_subclasses/fake_tensor.py:740 in infer_size)\n",
      "[2023-10-09 18:49:38,768] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s24, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:1405 in forward (_meta_registrations.py:3515 in common_meta_baddbmm_bmm)\n",
      "[2023-10-09 18:49:39,086] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval s0 + s25 <= 4096 [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:1476 in forward (_decomp/decompositions.py:725 in slice_forward)\n",
      "[2023-10-09 18:49:39,239] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s25, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:1515 in forward (_subclasses/fake_tensor.py:740 in infer_size)\n",
      "[2023-10-09 18:49:39,274] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s26, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:1522 in forward (_meta_registrations.py:3515 in common_meta_baddbmm_bmm)\n",
      "[2023-10-09 18:49:39,597] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval s0 + s27 <= 4096 [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:1593 in forward (_decomp/decompositions.py:725 in slice_forward)\n",
      "[2023-10-09 18:49:39,759] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s27, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:1632 in forward (_subclasses/fake_tensor.py:740 in infer_size)\n",
      "[2023-10-09 18:49:39,812] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s28, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:1639 in forward (_meta_registrations.py:3515 in common_meta_baddbmm_bmm)\n",
      "[2023-10-09 18:49:40,330] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval s0 + s29 <= 4096 [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:1710 in forward (_decomp/decompositions.py:725 in slice_forward)\n",
      "[2023-10-09 18:49:40,534] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s29, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:1749 in forward (_subclasses/fake_tensor.py:740 in infer_size)\n",
      "[2023-10-09 18:49:40,582] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s30, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:1756 in forward (_meta_registrations.py:3515 in common_meta_baddbmm_bmm)\n",
      "[2023-10-09 18:49:41,068] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval s0 + s31 <= 4096 [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:1827 in forward (_decomp/decompositions.py:725 in slice_forward)\n",
      "[2023-10-09 18:49:41,242] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s31, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:1866 in forward (_subclasses/fake_tensor.py:740 in infer_size)\n",
      "[2023-10-09 18:49:41,280] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s32, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:1873 in forward (_meta_registrations.py:3515 in common_meta_baddbmm_bmm)\n",
      "[2023-10-09 18:49:41,686] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval s0 + s33 <= 4096 [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:1944 in forward (_decomp/decompositions.py:725 in slice_forward)\n",
      "[2023-10-09 18:49:41,968] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s33, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:1983 in forward (_subclasses/fake_tensor.py:740 in infer_size)\n",
      "[2023-10-09 18:49:42,004] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s34, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:1990 in forward (_meta_registrations.py:3515 in common_meta_baddbmm_bmm)\n",
      "[2023-10-09 18:49:42,419] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval s0 + s35 <= 4096 [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:2061 in forward (_decomp/decompositions.py:725 in slice_forward)\n",
      "[2023-10-09 18:49:42,580] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s35, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:2100 in forward (_subclasses/fake_tensor.py:740 in infer_size)\n",
      "[2023-10-09 18:49:42,618] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s36, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:2107 in forward (_meta_registrations.py:3515 in common_meta_baddbmm_bmm)\n",
      "[2023-10-09 18:49:43,002] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval s0 + s37 <= 4096 [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:2178 in forward (_decomp/decompositions.py:725 in slice_forward)\n",
      "[2023-10-09 18:49:43,174] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s37, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:2217 in forward (_subclasses/fake_tensor.py:740 in infer_size)\n",
      "[2023-10-09 18:49:43,215] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s38, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:2224 in forward (_meta_registrations.py:3515 in common_meta_baddbmm_bmm)\n",
      "[2023-10-09 18:49:43,566] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval s0 + s39 <= 4096 [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:2295 in forward (_decomp/decompositions.py:725 in slice_forward)\n",
      "[2023-10-09 18:49:43,738] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s39, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:2334 in forward (_subclasses/fake_tensor.py:740 in infer_size)\n",
      "[2023-10-09 18:49:43,776] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s40, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:2341 in forward (_meta_registrations.py:3515 in common_meta_baddbmm_bmm)\n",
      "[2023-10-09 18:49:44,116] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval s0 + s41 <= 4096 [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:2412 in forward (_decomp/decompositions.py:725 in slice_forward)\n",
      "[2023-10-09 18:49:44,281] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s41, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:2451 in forward (_subclasses/fake_tensor.py:740 in infer_size)\n",
      "[2023-10-09 18:49:44,320] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s42, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:2458 in forward (_meta_registrations.py:3515 in common_meta_baddbmm_bmm)\n",
      "[2023-10-09 18:49:44,656] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval s0 + s43 <= 4096 [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:2529 in forward (_decomp/decompositions.py:725 in slice_forward)\n",
      "[2023-10-09 18:49:44,822] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s43, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:2568 in forward (_subclasses/fake_tensor.py:740 in infer_size)\n",
      "[2023-10-09 18:49:44,860] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s44, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:2575 in forward (_meta_registrations.py:3515 in common_meta_baddbmm_bmm)\n",
      "[2023-10-09 18:49:45,218] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval s0 + s45 <= 4096 [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:2646 in forward (_decomp/decompositions.py:725 in slice_forward)\n",
      "[2023-10-09 18:49:45,387] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s45, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:2685 in forward (_subclasses/fake_tensor.py:740 in infer_size)\n",
      "[2023-10-09 18:49:45,426] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s46, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:2692 in forward (_meta_registrations.py:3515 in common_meta_baddbmm_bmm)\n",
      "[2023-10-09 18:49:45,772] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval s0 + s47 <= 4096 [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:2763 in forward (_decomp/decompositions.py:725 in slice_forward)\n",
      "[2023-10-09 18:49:45,943] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s47, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:2802 in forward (_subclasses/fake_tensor.py:740 in infer_size)\n",
      "[2023-10-09 18:49:45,983] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s48, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:2809 in forward (_meta_registrations.py:3515 in common_meta_baddbmm_bmm)\n",
      "[2023-10-09 18:49:46,376] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval s0 + s49 <= 4096 [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:2880 in forward (_decomp/decompositions.py:725 in slice_forward)\n",
      "[2023-10-09 18:49:46,563] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s49, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:2919 in forward (_subclasses/fake_tensor.py:740 in infer_size)\n",
      "[2023-10-09 18:49:46,605] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s50, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:2926 in forward (_meta_registrations.py:3515 in common_meta_baddbmm_bmm)\n",
      "[2023-10-09 18:49:46,962] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval s0 + s51 <= 4096 [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:2997 in forward (_decomp/decompositions.py:725 in slice_forward)\n",
      "[2023-10-09 18:49:47,136] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s51, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:3036 in forward (_subclasses/fake_tensor.py:740 in infer_size)\n",
      "[2023-10-09 18:49:47,176] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s52, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:3043 in forward (_meta_registrations.py:3515 in common_meta_baddbmm_bmm)\n",
      "[2023-10-09 18:49:47,540] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval s0 + s53 <= 4096 [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:3114 in forward (_decomp/decompositions.py:725 in slice_forward)\n",
      "[2023-10-09 18:49:47,718] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s53, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:3153 in forward (_subclasses/fake_tensor.py:740 in infer_size)\n",
      "[2023-10-09 18:49:47,758] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s54, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:3160 in forward (_meta_registrations.py:3515 in common_meta_baddbmm_bmm)\n",
      "[2023-10-09 18:49:48,125] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval s0 + s55 <= 4096 [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:3231 in forward (_decomp/decompositions.py:725 in slice_forward)\n",
      "[2023-10-09 18:49:48,308] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s55, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:3270 in forward (_subclasses/fake_tensor.py:740 in infer_size)\n",
      "[2023-10-09 18:49:48,349] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s56, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:3277 in forward (_meta_registrations.py:3515 in common_meta_baddbmm_bmm)\n",
      "[2023-10-09 18:49:48,715] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval s0 + s57 <= 4096 [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:3348 in forward (_decomp/decompositions.py:725 in slice_forward)\n",
      "[2023-10-09 18:49:48,897] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s57, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:3387 in forward (_subclasses/fake_tensor.py:740 in infer_size)\n",
      "[2023-10-09 18:49:48,937] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s58, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:3394 in forward (_meta_registrations.py:3515 in common_meta_baddbmm_bmm)\n",
      "[2023-10-09 18:49:49,317] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval s0 + s59 <= 4096 [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:3465 in forward (_decomp/decompositions.py:725 in slice_forward)\n",
      "[2023-10-09 18:49:49,499] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s59, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:3504 in forward (_subclasses/fake_tensor.py:740 in infer_size)\n",
      "[2023-10-09 18:49:49,540] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s60, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:3511 in forward (_meta_registrations.py:3515 in common_meta_baddbmm_bmm)\n",
      "[2023-10-09 18:49:49,915] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval s0 + s61 <= 4096 [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:3582 in forward (_decomp/decompositions.py:725 in slice_forward)\n",
      "[2023-10-09 18:49:50,113] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s61, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:3621 in forward (_subclasses/fake_tensor.py:740 in infer_size)\n",
      "[2023-10-09 18:49:50,155] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s62, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:3628 in forward (_meta_registrations.py:3515 in common_meta_baddbmm_bmm)\n",
      "[2023-10-09 18:49:50,515] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval s0 + s63 <= 4096 [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:3699 in forward (_decomp/decompositions.py:725 in slice_forward)\n",
      "[2023-10-09 18:49:50,697] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s63, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:3738 in forward (_subclasses/fake_tensor.py:740 in infer_size)\n",
      "[2023-10-09 18:49:50,737] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0 + s64, s0 + s1) [guard added] at <eval_with_key>.5 from /home/dan/SHARK-Turbine/turbine_venv/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped:3745 in forward (_meta_registrations.py:3515 in common_meta_baddbmm_bmm)\n",
      "[2023-10-09 18:49:53,791] [1/0] torch.fx.experimental.symbolic_shapes: [INFO] produce_guards\n",
      "[2023-10-09 18:49:54,155] torch.fx.experimental.symbolic_shapes: [WARNING] Ignored guard s0 + s1 > 4096 == False, this could result in accuracy problems\n",
      "[2023-10-09 18:49:54,157] torch.fx.experimental.symbolic_shapes: [INFO] eval s0 + s1 <= 4096 [guard added] (_decomp/decompositions.py:725 in slice_forward)\n"
     ]
    }
   ],
   "source": [
    "#Run the export pipeline\n",
    "inst = StateUpdateModule(context=Context(), import_to=\"IMPORT\")\n",
    "module_str = str(CompiledModule.get_mlir_module(inst))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "bc04e1db-a8cc-4182-884d-ba3d8ae5adeb",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Output a torch-ir mlir file\n",
    "with open(\"llama2_torch.mlir\", \"w+\") as f:\n",
    "    f.write(module_str)\n",
    "#TODO: run the rest of the compile pipeline and do inference"
   ]
  }
 ],
 "metadata": {
  "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.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}