mnist.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from tinygrad import nn, Tensor\n",
    "from tinygrad.nn.optim import SGD\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Splitting training and test sets\n",
    "\n",
    "train_df = pd.read_csv(\"datasets/mnist-in-csv/mnist_train.csv\")\n",
    "test_df = pd.read_csv(\"datasets/mnist-in-csv/mnist_test.csv\")\n",
    "\n",
    "x_train = train_df.iloc[:, 1:].values # pixels\n",
    "y_train = train_df.iloc[:, 0].values # labels\n",
    "\n",
    "x_test = test_df.iloc[:, 1:].values # pixels\n",
    "y_test = test_df.iloc[:, 0].values # labels\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "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[4], line 42\u001b[0m\n\u001b[1;32m     40\u001b[0m net \u001b[38;5;241m=\u001b[39m Network()\n\u001b[1;32m     41\u001b[0m \u001b[38;5;66;03m#print(x_train[].shape)\u001b[39;00m\n\u001b[0;32m---> 42\u001b[0m \u001b[43mnet\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtrain\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mx_train\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43my_train\u001b[49m\u001b[43m)\u001b[49m\n",
      "Cell \u001b[0;32mIn[4], line 33\u001b[0m, in \u001b[0;36mNetwork.train\u001b[0;34m(self, epochs, data, labels)\u001b[0m\n\u001b[1;32m     31\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mopt\u001b[38;5;241m.\u001b[39mzero_grad()\n\u001b[1;32m     32\u001b[0m loss\u001b[38;5;241m.\u001b[39mbackward()\n\u001b[0;32m---> 33\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mopt\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     34\u001b[0m Y \u001b[38;5;241m=\u001b[39m pred\u001b[38;5;241m.\u001b[39margmax(axis\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m)\n\u001b[1;32m     35\u001b[0m acc \u001b[38;5;241m=\u001b[39m (Y \u001b[38;5;241m==\u001b[39m labels)\u001b[38;5;241m.\u001b[39mmean()\n",
      "File \u001b[0;32m/opt/homebrew/lib/python3.11/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/homebrew/lib/python3.11/site-packages/tinygrad/tensor.py:3256\u001b[0m, in \u001b[0;36m_metadata_wrapper.<locals>._wrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m   3253\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m: caller \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m   3255\u001b[0m token \u001b[38;5;241m=\u001b[39m _METADATA\u001b[38;5;241m.\u001b[39mset(Metadata(name\u001b[38;5;241m=\u001b[39mfn\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m, caller\u001b[38;5;241m=\u001b[39mcaller))\n\u001b[0;32m-> 3256\u001b[0m ret \u001b[38;5;241m=\u001b[39m \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   3257\u001b[0m _METADATA\u001b[38;5;241m.\u001b[39mreset(token)\n\u001b[1;32m   3258\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m ret\n",
      "File \u001b[0;32m/opt/homebrew/lib/python3.11/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/homebrew/lib/python3.11/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/homebrew/lib/python3.11/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/homebrew/lib/python3.11/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/homebrew/lib/python3.11/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/homebrew/lib/python3.11/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/homebrew/lib/python3.11/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/homebrew/lib/python3.11/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": [
    "def activation(x: Tensor) -> Tensor:\n",
    "    return x.abs()\n",
    "\n",
    "def sparse_categorical_crossentropy(self, Y, ignore_index=-1) -> Tensor:\n",
    "    loss_mask = Y != ignore_index\n",
    "    y_counter = Tensor.arange(self.shape[-1], requires_grad=False, device=self.device).unsqueeze(0).expand(Y.numel(), self.shape[-1])\n",
    "    y = ((y_counter == Y.flatten().reshape(-1, 1)).where(-1.0, 0) * loss_mask.reshape(-1, 1)).reshape(*Y.shape, self.shape[-1])\n",
    "    return self.log_softmax().mul(y).sum() / loss_mask.sum()\n",
    "\n",
    "class Network:\n",
    "    def __init__(self):\n",
    "        self.layers = [nn.Linear(784, 128, bias=False), nn.Linear(128, 10, bias=False)]\n",
    "        self.opt = SGD([self.layers[0].weight, self.layers[1].weight], lr=3e-4)\n",
    "\n",
    "    def forward(self, x: Tensor):\n",
    "        x = self.layers[0](x)\n",
    "        x = activation(x)\n",
    "        x = self.layers[1](x)\n",
    "\n",
    "        return x\n",
    "\n",
    "    def train(self, epochs: int, data: Tensor, labels: Tensor):\n",
    "        with Tensor.train():\n",
    "            for i in range(epochs):\n",
    "                # sample data\n",
    "                sample = np.random.randint(0, x_train.shape[0], size=64)\n",
    "                batch = Tensor(x_train[sample], requires_grad=False)\n",
    "                labels = Tensor(y_train[sample])\n",
    "                pred = self.forward(batch)\n",
    "                loss = sparse_categorical_crossentropy(pred, labels)\n",
    "                self.opt.zero_grad()\n",
    "                loss.backward()\n",
    "                self.opt.step()\n",
    "                Y = pred.argmax(axis=-1)\n",
    "                acc = (Y == labels).mean()\n",
    "                if i%100 == 0:\n",
    "                    print(f\"Step {i+1} | Loss: {loss.numpy()} | Accuracy: {acc.numpy()}\")\n",
    "\n",
    "\n",
    "net = Network()\n",
    "#print(x_train[].shape)\n",
    "net.train(1, x_train, y_train)\n",
    "\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
}