playground.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = torch.arange(1, 13, 1).reshape(3, 4).float()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 1.,  2.,  3.,  4.],\n",
       "        [ 5.,  6.,  7.,  8.],\n",
       "        [ 9., 10., 11., 12.]])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "a.requires_grad = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = torch.arange(1, 7, 1).reshape(2, 3).float()\n",
    "a.requires_grad = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.118033988749895\n",
      "2.5\n",
      "[[2.25 0.25]\n",
      " [0.25 2.25]]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([[1.11803399]])"
      ]
     },
     "execution_count": 62,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a = np.array([[1, 2], [3, 4]])\n",
    "\n",
    "std = np.std(a, keepdims=False)\n",
    "print(std)\n",
    "\n",
    "temp = (a - a.mean(keepdims=True)) ** 2\n",
    "\n",
    "print(a.mean())\n",
    "print(temp)\n",
    "\n",
    "std = np.sqrt(temp.mean(keepdims=True))\n",
    "std"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0.5000, 0.5000], grad_fn=<VarBackward0>)\n",
      "tensor([[-1.,  1.],\n",
      "        [-1.,  1.]])\n"
     ]
    }
   ],
   "source": [
    "a = torch.Tensor([[1, 2], [3, 4]])\n",
    "a = a.to(torch.float32)\n",
    "a.requires_grad = True\n",
    "\n",
    "b = torch.var(a, axis=1)\n",
    "c = b.sum()\n",
    "\n",
    "c.backward()\n",
    "\n",
    "print(b)\n",
    "print(a.grad)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output of the loss function:  tensor(2., grad_fn=<MseLossBackward0>)\n",
      "tensor([ 0., -2.,  2.,  0.,  0.,  0.])\n"
     ]
    }
   ],
   "source": [
    "pred = torch.arange(1, 7, 1).float(); pred.requires_grad = True\n",
    "actual = torch.Tensor([1, 3, 2, 4, 5, 6])\n",
    "\n",
    "loss = nn.MSELoss(reduction=\"sum\")\n",
    "\n",
    "output: torch.Tensor = loss(pred, actual)\n",
    "output.backward()\n",
    "\n",
    "print(\"Output of the loss function: \", output)\n",
    "print(pred.grad)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[1., 2.],\n",
       "         [3., 4.]], requires_grad=True),\n",
       " tensor([[1., 2.],\n",
       "         [3., 4.]], dtype=torch.float16, grad_fn=<ToCopyBackward0>),\n",
       " tensor([[1., 2.],\n",
       "         [3., 4.]], dtype=torch.float16, grad_fn=<ToCopyBackward0>),\n",
       " tensor([[2., 4.],\n",
       "         [6., 8.]], dtype=torch.float16, grad_fn=<AddBackward0>),\n",
       " tensor(20., dtype=torch.float16, grad_fn=<SumBackward0>))"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "d = a.half()\n",
    "k = a.half()\n",
    "\n",
    "e = d + k\n",
    "f = e.sum()\n",
    "\n",
    "a, d, k, e, f"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "f.backward()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[2.5000, 2.5000],\n",
       "        [2.5000, 2.5000]])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a.grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[1., 2.],\n",
      "        [3., 4.]], requires_grad=True) tensor(10., grad_fn=<SumBackward0>)\n"
     ]
    }
   ],
   "source": [
    "x = torch.tensor([[1., 2.], [3., 4.]], requires_grad=True)\n",
    "\n",
    "y = x.sum()  # y is a vector: tensor([4., 6.], grad_fn=<SumBackward1>)\n",
    "\n",
    "print(x, y)\n",
    "y.retain_grad = True\n",
    "y.backward()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1., 1.],\n",
       "        [1., 1.]])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x.grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = np.arange(1, 13, 1).reshape(3, 4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 2,  3,  4,  5],\n",
       "       [ 6,  7,  8,  9],\n",
       "       [10, 11, 12, 13]])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b = a + 1\n",
    "b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output from Linear layer:  tensor([ 2.1948,  0.2230,  1.8296,  1.1871,  1.2304,  1.1347, -1.8223,  2.1204,\n",
      "        -0.1979, -3.8751], grad_fn=<ViewBackward0>)\n",
      "Input shape:  torch.Size([5])\n",
      "Output shape:  torch.Size([10])\n",
      "torch.Size([10, 5]) torch.Size([10])\n"
     ]
    }
   ],
   "source": [
    "import torch.nn as nn\n",
    "\n",
    "layer = nn.Linear(5, 10)\n",
    "\n",
    "inp = torch.Tensor([1, 2, 3, 4, 5])\n",
    "\n",
    "out = layer(inp)\n",
    "\n",
    "print(\"Output from Linear layer: \", out)\n",
    "print(\"Input shape: \", inp.shape)\n",
    "print(\"Output shape: \", out.shape)\n",
    "\n",
    "print(layer.weight.shape, layer.bias.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([3., 4., 1., 2.])\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "# Create a tensor with requires_grad=True\n",
    "x = torch.tensor([1.0, 2.0, 3.0, 4.0], requires_grad=True)\n",
    "\n",
    "# Split the tensor into two chunks\n",
    "a, b = torch.split(x, 2)\n",
    "\n",
    "# Perform some operations on the split tensors\n",
    "y = (a * b).sum()\n",
    "\n",
    "# Compute the gradients\n",
    "y.backward()\n",
    "\n",
    "# Print the gradients\n",
    "print(x.grad)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[[ 1  2]\n",
      "  [ 3  4]\n",
      "  [ 5  6]]\n",
      "\n",
      " [[ 7  8]\n",
      "  [ 9 10]\n",
      "  [11 12]]]\n",
      "--------\n",
      "[functools.partial(<function back at 0x7fdb9869cdc0>, index=0), functools.partial(<function back at 0x7fdb9869cdc0>, index=1)]\n",
      "after  1\n",
      "[[[ 2  2]\n",
      "  [ 6  4]\n",
      "  [10  6]]\n",
      "\n",
      " [[14  8]\n",
      "  [18 10]\n",
      "  [22 12]]]\n",
      "after  2\n",
      "[[[ 2  4]\n",
      "  [ 6  8]\n",
      "  [10 12]]\n",
      "\n",
      " [[14 16]\n",
      "  [18 20]\n",
      "  [22 24]]]\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "# Create a random array 'x' with a random shape\n",
    "x = np.arange(1, 13, 1).reshape((2, 3, 2))\n",
    "print(x)\n",
    "\n",
    "# Choose a random axis 'n' to split along\n",
    "n = -1\n",
    "\n",
    "# Choose a random number of splits 'k'\n",
    "k = 2\n",
    "\n",
    "# Split 'x' into 'k' parts along axis 'n'\n",
    "parts = np.split(x, k, axis=n)\n",
    "\n",
    "# Create index arrays for each part\n",
    "indices = []\n",
    "start = 0\n",
    "for part in parts:\n",
    "    idx = [slice(None)] * x.ndim\n",
    "    idx[n] = slice(start, start + part.shape[n])\n",
    "    indices.append(tuple(idx))\n",
    "    start += part.shape[n]\n",
    "\n",
    "\n",
    "from functools import partial\n",
    "\n",
    "def back(index=0):\n",
    "    part, idx = parts[index], indices[index]\n",
    "    x[idx] += part\n",
    "\n",
    "\n",
    "funcs = []\n",
    "for i in range(len(parts)):\n",
    "    f = partial(back, index = i)\n",
    "    funcs.append(f)\n",
    "\n",
    "print(\"--------\")\n",
    "print(funcs)\n",
    "\n",
    "# Update the corresponding portions of 'x' with each part\n",
    "for i, f in enumerate(funcs):\n",
    "    f()\n",
    "    print(\"after \", i+1)\n",
    "    print(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 4, 9])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "q = torch.arange(1, 25, 1).reshape(2, 4, 3)\n",
    "k = torch.arange(1, 25, 1).reshape(2, 4, 3)\n",
    "v = torch.arange(1, 25, 1).reshape(2, 4, 3)\n",
    "\n",
    "x = None\n",
    "x = torch.cat([q, k], dim=-1)\n",
    "x = torch.cat([x, v], dim=-1)\n",
    "\n",
    "x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 4, 9])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y = torch.cat([q, k, v], dim=-1)\n",
    "y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "l1 = [q]\n",
    "l2 = [k, v]\n",
    "\n",
    "final = l1 + l2\n",
    "len(final)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "q = np.arange(1, 25, 1).reshape(2, 4, 3)\n",
    "k = np.arange(1, 9, 1).reshape(2, 4, 1)\n",
    "v = np.arange(1, 17, 1).reshape(2, 4, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2, 4, 6)"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dim = -1\n",
    "\n",
    "x = np.concatenate([q, k, v], axis=dim)\n",
    "x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([3, 4])"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "l = [q, k, v]\n",
    "sizes = [hi.shape[dim] for hi in l]\n",
    "\n",
    "cum = np.cumsum(sizes[:-1])\n",
    "cum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "True\n",
      "True\n"
     ]
    }
   ],
   "source": [
    "ys = np.split(x, cum, axis=dim)\n",
    "\n",
    "print(np.array_equal(ys[0], q))\n",
    "print(np.array_equal(ys[1], k))\n",
    "print(np.array_equal(ys[2], v))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2 5 6]\n",
      "[2, 5, 6]\n"
     ]
    }
   ],
   "source": [
    "import mlx.core as mx\n",
    "import numpy as np\n",
    "\n",
    "a = [2, 3, 1]\n",
    "b = mx.array(a)\n",
    "c = np.array(a)\n",
    "\n",
    "print(np.cumsum(b))\n",
    "\n",
    "print(mx.cumsum(b).tolist())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "all_false = mx.all(b)\n",
    "all_false.item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "c = np.array(a)\n",
    "all_false = np.all(c)\n",
    "all_false.item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = np.arange(1, 25, 1).reshape(2, 2, 6)\n",
    "heads = 3\n",
    "b = np.reshape(a, (2, heads, 2, 6 // heads))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2, 2, 6)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b.reshape(2, 2, 6).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 3, 4])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "a = torch.arange(1, 25, 1).reshape(2, 3, 4).float()\n",
    "a.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "drop = torch.nn.Dropout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 0.,  4.,  6.,  0.],\n",
       "         [ 0.,  0., 14., 16.],\n",
       "         [ 0., 20.,  0.,  0.]],\n",
       "\n",
       "        [[26., 28., 30.,  0.],\n",
       "         [ 0., 36.,  0., 40.],\n",
       "         [42.,  0., 46., 48.]]])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b = drop(a)\n",
    "b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 1.,  2.,  3.,  4.],\n",
       "         [ 5.,  6.,  7.,  8.],\n",
       "         [ 9., 10., 11., 12.]],\n",
       "\n",
       "        [[13., 14., 15., 16.],\n",
       "         [17., 18., 19., 20.],\n",
       "         [21., 22., 23., 24.]]])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "drop.eval()\n",
    "c = drop(a)\n",
    "c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mlx.core as mx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([[1, 6, 3],\n",
       "        [5, 3, 2]], dtype=float32),\n",
       " array([[0, 1, 0],\n",
       "        [1, 0, 0]], dtype=int8))"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a = mx.array([[1, 6, 3], [5, 3, 2]], dtype=mx.float32)\n",
    "mask = mx.array([[0, 1, 0], [1, 0, 0]], dtype=mx.int8)\n",
    "\n",
    "a, mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[1, -1, 3],\n",
       "       [-1, 3, 2]], dtype=float32)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b = mx.where(mask, mx.array(-1), a)\n",
    "b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([[1, 6, 3],\n",
       "        [5, 3, 2]], dtype=float32),\n",
       " array([[0, 1, 0],\n",
       "        [1, 0, 0]], dtype=int8))"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "a = mx.array([[1, 6, 3], [5, 3, 2]], dtype=mx.float32)\n",
    "mask = mx.array([[0, 1, 0], [1, 0, 0]], dtype=mx.int8)\n",
    "\n",
    "a, mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 1., -1.,  3.],\n",
       "       [-1.,  3.,  2.]], dtype=float32)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b = np.where(mask, np.array(-1), a)\n",
    "b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}