transformer.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 167,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tinygrad import Tensor, nn\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 168,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hyperparams\n",
    "\n",
    "training_steps = 1000\n",
    "train_split = 0.9 # 90% of the data goes to training, 10% for validation\n",
    "block_size = 8 # length of sequence\n",
    "batch_size = 4 # number of sequences\n",
    "vocab_size = 28\n",
    "head_size = 16 # size of each attention head\n",
    "learning_rate = 1e-4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 169,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([17, 14, 15, 25,  2,  5, 27, 22, 23, 25, 19,  7, 26,  0, 24, 25, 11,\n",
       "       26,  6, 25, 10,  5, 18,  8,  9, 25, 26, 13, 15,  7, 25, 21, 14, 15,\n",
       "       25,  4, 20, 12, 16, 25,  3, 26,  1])"
      ]
     },
     "execution_count": 169,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def tokenize(input: str) -> np.ndarray:\n",
    "    \"\"\"\n",
    "    Tokenizes an input string.\n",
    "    \"\"\"\n",
    "    vocab = set(input)\n",
    "    mapping = {v:i for i, v in enumerate(list(vocab))}\n",
    "    return np.array([mapping[char] for char in input])\n",
    "\n",
    "tokenize(\"The quick brown fox jumps over the lazy dog\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 170,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"datasets/tiny_shakespeare.txt\") as f:\n",
    "    input_text = f.read()\n",
    "    data = tokenize(input_text)\n",
    "    n = int(train_split * len(data))\n",
    "    train_data = data[:n]\n",
    "    test_data = data[n:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 171,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(<Tensor <LB METAL (4, 8) long (<MetaOps.COPY: 3>, None)> on METAL with grad None>,\n",
       " <Tensor <LB METAL (4, 8) long (<MetaOps.COPY: 3>, None)> on METAL with grad None>)"
      ]
     },
     "execution_count": 171,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def get_batch(split: str) -> tuple[Tensor, Tensor]:\n",
    "    \"\"\"\n",
    "    Given input data of shape (N,), return a randomly sampled batch of shape (batch_size, block_size)\n",
    "    \"\"\"\n",
    "    split_data = train_data if split == \"train\" else test_data\n",
    "    sample = np.random.randint(0, split_data.shape[0]-block_size, size=(batch_size))\n",
    "    x = np.stack([split_data[i:i+block_size] for i in sample])\n",
    "    y = np.stack([split_data[i+1:i+block_size+1] for i in sample])\n",
    "    return Tensor(x), Tensor(y)\n",
    "\n",
    "get_batch(\"train\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 172,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SelfAttentionHead:\n",
    "    def __init__(self):\n",
    "        self.queries = nn.Linear(vocab_size, head_size, bias=False)\n",
    "        self.keys = nn.Linear(vocab_size, head_size, bias=False)\n",
    "        self.values = nn.Linear(vocab_size, head_size, bias=False)\n",
    "        self.projection = nn.Linear(head_size, vocab_size, bias=False)  # Add projection layer\n",
    "    \n",
    "    def __call__(self, input):\n",
    "        # input: (B, T, C)\n",
    "        q = self.queries(input)  # (B, T, 16)\n",
    "        k = self.keys(input)      # (B, T, 16)\n",
    "        v = self.values(input)    # (B, T, 16)\n",
    "        x = q @ k.transpose(-2, -1) * head_size**-0.5  # (B, T, T)\n",
    "        tril = Tensor.tril(Tensor.ones(block_size, block_size))\n",
    "        x = x.masked_fill(tril == 0, float('-inf'))\n",
    "        x = Tensor.softmax(x, axis=-1) @ v  # (B, T, 16)\n",
    "\n",
    "        x = self.projection(x)  # Project output to vocab_size -> (B, T, vocab_size)\n",
    "        return x\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 173,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FeedForwardNetwork:\n",
    "    def __init__(self):\n",
    "        self.l1 = nn.Linear(vocab_size, head_size)\n",
    "        self.l2 = nn.Linear(head_size, vocab_size)\n",
    "    def __call__(self, x):\n",
    "        x = self.l1(x)\n",
    "        x = Tensor.relu(x)\n",
    "        x = self.l2(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 174,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Model:\n",
    "    def __init__(self):\n",
    "        self.token_embeddings = nn.Embedding(vocab_size, vocab_size)\n",
    "        self.positional_embeddings = nn.Embedding(block_size, vocab_size)\n",
    "        self.attention = SelfAttentionHead()\n",
    "        self.ff = FeedForwardNetwork()\n",
    "        self.ln = nn.LayerNorm(vocab_size)\n",
    "        self.lm_head = nn.Linear(vocab_size, vocab_size)\n",
    "\n",
    "    def __call__(self, idx, targets=None):\n",
    "        token_embedding = self.token_embeddings(idx)  # (B, T, C)\n",
    "        pos_embedding = self.positional_embeddings(Tensor.arange(block_size))  # (T, C)\n",
    "        x = token_embedding + pos_embedding  # (B, T, C)\n",
    "        x = self.ln(x)  # Add normalization before attention\n",
    "        x = x + self.attention(x)\n",
    "        x = self.ln(x)  # Add normalization after attention\n",
    "        x = x + self.ff(x)\n",
    "        logits = self.lm_head(x)\n",
    "        B, T, C = logits.shape\n",
    "        logits = logits.view(B*T, C)\n",
    "        targets = targets.view(B*T)\n",
    "        loss = Tensor.sparse_categorical_crossentropy(logits, targets)\n",
    "        return logits, loss\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 175,
   "metadata": {},
   "outputs": [
    {
     "ename": "AssertionError",
     "evalue": "Invalid Metal library. Could be due to using conda. Try system python or METAL_XCODE=1 DISABLE_COMPILER_CACHE=1.",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAssertionError\u001b[0m                            Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[175], line 9\u001b[0m\n\u001b[1;32m      7\u001b[0m optim\u001b[38;5;241m.\u001b[39mzero_grad()\n\u001b[1;32m      8\u001b[0m loss\u001b[38;5;241m.\u001b[39mbackward()\n\u001b[0;32m----> 9\u001b[0m \u001b[43moptim\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstep\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     11\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m step \u001b[38;5;241m%\u001b[39m \u001b[38;5;241m100\u001b[39m \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m0\u001b[39m:\n\u001b[1;32m     12\u001b[0m     \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mStep \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mstep\u001b[38;5;241m+\u001b[39m\u001b[38;5;241m1\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m | Loss: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mloss\u001b[38;5;241m.\u001b[39mnumpy()\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n",
      "File \u001b[0;32m/opt/anaconda3/envs/gym/lib/python3.12/site-packages/tinygrad/nn/optim.py:34\u001b[0m, in \u001b[0;36mOptimizer.step\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m     30\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mstep\u001b[39m(\u001b[38;5;28mself\u001b[39m):\n\u001b[1;32m     31\u001b[0m \u001b[38;5;250m  \u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m     32\u001b[0m \u001b[38;5;124;03m  Performs a single optimization step.\u001b[39;00m\n\u001b[1;32m     33\u001b[0m \u001b[38;5;124;03m  \"\"\"\u001b[39;00m\n\u001b[0;32m---> 34\u001b[0m   \u001b[43mTensor\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrealize\u001b[49m\u001b[43m(\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[43mschedule_step\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/opt/anaconda3/envs/gym/lib/python3.12/site-packages/tinygrad/tensor.py:3231\u001b[0m, in \u001b[0;36m_metadata_wrapper.<locals>._wrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m   3230\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_wrapper\u001b[39m(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs):\n\u001b[0;32m-> 3231\u001b[0m   \u001b[38;5;28;01mif\u001b[39;00m _METADATA\u001b[38;5;241m.\u001b[39mget() \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;01mreturn\u001b[39;00m \u001b[43mfn\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\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[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   3233\u001b[0m   \u001b[38;5;28;01mif\u001b[39;00m TRACEMETA \u001b[38;5;241m>\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m2\u001b[39m:\n\u001b[1;32m   3234\u001b[0m     caller_frame \u001b[38;5;241m=\u001b[39m sys\u001b[38;5;241m.\u001b[39m_getframe(frame \u001b[38;5;241m:=\u001b[39m \u001b[38;5;241m1\u001b[39m)\n",
      "File \u001b[0;32m/opt/anaconda3/envs/gym/lib/python3.12/site-packages/tinygrad/tensor.py:204\u001b[0m, in \u001b[0;36mTensor.realize\u001b[0;34m(self, do_update_stats, *lst)\u001b[0m\n\u001b[1;32m    202\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mrealize\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;241m*\u001b[39mlst:Tensor, do_update_stats\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Tensor:\n\u001b[1;32m    203\u001b[0m \u001b[38;5;250m  \u001b[39m\u001b[38;5;124;03m\"\"\"Triggers the computation needed to create these Tensor(s).\"\"\"\u001b[39;00m\n\u001b[0;32m--> 204\u001b[0m   \u001b[43mrun_schedule\u001b[49m\u001b[43m(\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[43mschedule_with_vars\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mlst\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdo_update_stats\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdo_update_stats\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    205\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\n",
      "File \u001b[0;32m/opt/anaconda3/envs/gym/lib/python3.12/site-packages/tinygrad/engine/realize.py:221\u001b[0m, in \u001b[0;36mrun_schedule\u001b[0;34m(schedule, var_vals, do_update_stats)\u001b[0m\n\u001b[1;32m    220\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mrun_schedule\u001b[39m(schedule:List[ScheduleItem], var_vals:Optional[Dict[Variable, \u001b[38;5;28mint\u001b[39m]]\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m, do_update_stats\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m):\n\u001b[0;32m--> 221\u001b[0m \u001b[43m  \u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mei\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mlower_schedule\u001b[49m\u001b[43m(\u001b[49m\u001b[43mschedule\u001b[49m\u001b[43m)\u001b[49m\u001b[43m:\u001b[49m\n\u001b[1;32m    222\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43;01mif\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43mlen\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mcapturing\u001b[49m\u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mand\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mCAPTURING\u001b[49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mcapturing\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m]\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43madd\u001b[49m\u001b[43m(\u001b[49m\u001b[43mei\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    223\u001b[0m \u001b[43m    \u001b[49m\u001b[43mei\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun\u001b[49m\u001b[43m(\u001b[49m\u001b[43mvar_vals\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdo_update_stats\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdo_update_stats\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/opt/anaconda3/envs/gym/lib/python3.12/site-packages/tinygrad/engine/realize.py:214\u001b[0m, in \u001b[0;36mlower_schedule\u001b[0;34m(schedule)\u001b[0m\n\u001b[1;32m    212\u001b[0m   \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtensor operations:\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m    213\u001b[0m   pprint\u001b[38;5;241m.\u001b[39mpprint(si\u001b[38;5;241m.\u001b[39mmetadata, indent\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m2\u001b[39m)\n\u001b[0;32m--> 214\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m e\n",
      "File \u001b[0;32m/opt/anaconda3/envs/gym/lib/python3.12/site-packages/tinygrad/engine/realize.py:208\u001b[0m, in \u001b[0;36mlower_schedule\u001b[0;34m(schedule)\u001b[0m\n\u001b[1;32m    206\u001b[0m \u001b[38;5;28;01mwhile\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(schedule):\n\u001b[1;32m    207\u001b[0m   si \u001b[38;5;241m=\u001b[39m schedule\u001b[38;5;241m.\u001b[39mpop(\u001b[38;5;241m0\u001b[39m)\n\u001b[0;32m--> 208\u001b[0m   \u001b[38;5;28;01mtry\u001b[39;00m: \u001b[38;5;28;01myield\u001b[39;00m \u001b[43mlower_schedule_item\u001b[49m\u001b[43m(\u001b[49m\u001b[43msi\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    209\u001b[0m   \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[1;32m    210\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m DEBUG \u001b[38;5;241m>\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m2\u001b[39m:\n",
      "File \u001b[0;32m/opt/anaconda3/envs/gym/lib/python3.12/site-packages/tinygrad/engine/realize.py:192\u001b[0m, in \u001b[0;36mlower_schedule_item\u001b[0;34m(si)\u001b[0m\n\u001b[1;32m    190\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(\u001b[38;5;28mset\u001b[39m(x\u001b[38;5;241m.\u001b[39mdevice \u001b[38;5;28;01mfor\u001b[39;00m x \u001b[38;5;129;01min\u001b[39;00m si\u001b[38;5;241m.\u001b[39mbufs)) \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m1\u001b[39m \u001b[38;5;129;01mor\u001b[39;00m (si\u001b[38;5;241m.\u001b[39mast\u001b[38;5;241m.\u001b[39mop \u001b[38;5;129;01mis\u001b[39;00m MetaOps\u001b[38;5;241m.\u001b[39mEXT \u001b[38;5;129;01mand\u001b[39;00m si\u001b[38;5;241m.\u001b[39mast\u001b[38;5;241m.\u001b[39marg[\u001b[38;5;241m0\u001b[39m] \u001b[38;5;129;01mis\u001b[39;00m MetaOps\u001b[38;5;241m.\u001b[39mCOPY) \u001b[38;5;129;01mor\u001b[39;00m getenv(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mUSE_COPY_KERNEL\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m    191\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m si\u001b[38;5;241m.\u001b[39mast\u001b[38;5;241m.\u001b[39mop \u001b[38;5;129;01mis\u001b[39;00m MetaOps\u001b[38;5;241m.\u001b[39mKERNEL:\n\u001b[0;32m--> 192\u001b[0m   runner \u001b[38;5;241m=\u001b[39m \u001b[43mget_runner\u001b[49m\u001b[43m(\u001b[49m\u001b[43msi\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43moutputs\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m]\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdevice\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msi\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mast\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    193\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m ExecItem(runner, [si\u001b[38;5;241m.\u001b[39mbufs[x] \u001b[38;5;28;01mfor\u001b[39;00m x \u001b[38;5;129;01min\u001b[39;00m runner\u001b[38;5;241m.\u001b[39mp\u001b[38;5;241m.\u001b[39mglobals], si\u001b[38;5;241m.\u001b[39mmetadata)\n\u001b[1;32m    194\u001b[0m out, (op, arg) \u001b[38;5;241m=\u001b[39m si\u001b[38;5;241m.\u001b[39moutputs[\u001b[38;5;241m0\u001b[39m], si\u001b[38;5;241m.\u001b[39mast\u001b[38;5;241m.\u001b[39marg\n",
      "File \u001b[0;32m/opt/anaconda3/envs/gym/lib/python3.12/site-packages/tinygrad/engine/realize.py:161\u001b[0m, in \u001b[0;36mget_runner\u001b[0;34m(dname, ast)\u001b[0m\n\u001b[1;32m    159\u001b[0m     \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtest\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mexternal\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mfuzz_uops\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m UOpsFuzzerRunner\n\u001b[1;32m    160\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m UOpsFuzzerRunner(replace(prg, dname\u001b[38;5;241m=\u001b[39mdname))\n\u001b[0;32m--> 161\u001b[0m   method_cache[ckey] \u001b[38;5;241m=\u001b[39m method_cache[bkey] \u001b[38;5;241m=\u001b[39m ret \u001b[38;5;241m=\u001b[39m \u001b[43mCompiledRunner\u001b[49m\u001b[43m(\u001b[49m\u001b[43mreplace\u001b[49m\u001b[43m(\u001b[49m\u001b[43mprg\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdname\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdname\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    162\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m ret\n",
      "File \u001b[0;32m/opt/anaconda3/envs/gym/lib/python3.12/site-packages/tinygrad/engine/realize.py:84\u001b[0m, in \u001b[0;36mCompiledRunner.__init__\u001b[0;34m(self, p, precompiled)\u001b[0m\n\u001b[1;32m     82\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mp:Program \u001b[38;5;241m=\u001b[39m p\n\u001b[1;32m     83\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlib:\u001b[38;5;28mbytes\u001b[39m \u001b[38;5;241m=\u001b[39m precompiled \u001b[38;5;28;01mif\u001b[39;00m precompiled \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 Device[p\u001b[38;5;241m.\u001b[39mdname]\u001b[38;5;241m.\u001b[39mcompiler\u001b[38;5;241m.\u001b[39mcompile_cached(p\u001b[38;5;241m.\u001b[39msrc)\n\u001b[0;32m---> 84\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mclprg \u001b[38;5;241m=\u001b[39m \u001b[43mDevice\u001b[49m\u001b[43m[\u001b[49m\u001b[43mp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdname\u001b[49m\u001b[43m]\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mruntime\u001b[49m\u001b[43m(\u001b[49m\u001b[43mp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfunction_name\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlib\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     85\u001b[0m \u001b[38;5;28msuper\u001b[39m()\u001b[38;5;241m.\u001b[39m\u001b[38;5;21m__init__\u001b[39m(p\u001b[38;5;241m.\u001b[39mname, p\u001b[38;5;241m.\u001b[39mdname, p\u001b[38;5;241m.\u001b[39mop_estimate, p\u001b[38;5;241m.\u001b[39mmem_estimate, p\u001b[38;5;241m.\u001b[39mlds_estimate)\n",
      "File \u001b[0;32m/opt/anaconda3/envs/gym/lib/python3.12/site-packages/tinygrad/runtime/ops_metal.py:39\u001b[0m, in \u001b[0;36mMetalProgram.__init__\u001b[0;34m(self, device, name, lib)\u001b[0m\n\u001b[1;32m     37\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m ret:\n\u001b[1;32m     38\u001b[0m       \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mError running disassembler: Make sure you have https://github.com/dougallj/applegpu cloned to tinygrad/extra/disassemblers/applegpu\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m---> 39\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m lib[:\u001b[38;5;241m4\u001b[39m] \u001b[38;5;241m==\u001b[39m \u001b[38;5;124mb\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mMTLB\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mInvalid Metal library. Could be due to using conda. Try system python or METAL_XCODE=1 DISABLE_COMPILER_CACHE=1.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m     40\u001b[0m data \u001b[38;5;241m=\u001b[39m libdispatch\u001b[38;5;241m.\u001b[39mdispatch_data_create(lib, \u001b[38;5;28mlen\u001b[39m(lib), \u001b[38;5;28;01mNone\u001b[39;00m, \u001b[38;5;28;01mNone\u001b[39;00m)\n\u001b[1;32m     41\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlibrary \u001b[38;5;241m=\u001b[39m unwrap2(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdevice\u001b[38;5;241m.\u001b[39mdevice\u001b[38;5;241m.\u001b[39mnewLibraryWithData_error_(data, \u001b[38;5;28;01mNone\u001b[39;00m))\n",
      "\u001b[0;31mAssertionError\u001b[0m: Invalid Metal library. Could be due to using conda. Try system python or METAL_XCODE=1 DISABLE_COMPILER_CACHE=1."
     ]
    }
   ],
   "source": [
    "m = Model()\n",
    "optim = nn.optim.AdamW(nn.state.get_parameters(m), lr=learning_rate)\n",
    "with Tensor.train():\n",
    "    for step in range(training_steps):\n",
    "        xb, yb = get_batch(\"train\") # (B, T)\n",
    "        _, loss = m(xb, yb)\n",
    "        optim.zero_grad()\n",
    "        loss.backward()\n",
    "        optim.step()\n",
    "\n",
    "        if step % 100 == 0:\n",
    "            print(f\"Step {step+1} | Loss: {loss.numpy()}\")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "gym",
   "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}