{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "691e4daf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Tensor(data=[2.], grad=None), Tensor(data=[3.], grad=None))"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "\n",
    "class Tensor:\n",
    "    def __init__(self, data, requires_grad=False):\n",
    "        self.data = np.asarray(data, dtype=float)\n",
    "        self.requires_grad = requires_grad\n",
    "        self.grad = None\n",
    "        self._op = None\n",
    "        self._parents = ()\n",
    "    def _accumulate_grad(self, grad):\n",
    "        if self.grad is None:\n",
    "            self.grad = np.zeros_like(self.data)\n",
    "        self.grad += grad\n",
    "\n",
    "    def backward(self):\n",
    "        order = topo_sort(self)\n",
    "        self.grad = np.ones_like(self.data)\n",
    "        for node in reversed(order):\n",
    "            if node._op == \"mul\":\n",
    "                left, right = node._parents\n",
    "                left._accumulate_grad(node.grad * right.data)\n",
    "                right._accumulate_grad(node.grad * left.data)\n",
    "\n",
    "            elif node._op == \"neg\":\n",
    "                parent, = node._parents\n",
    "                parent._accumulate_grad(-node.grad)\n",
    "\n",
    "            elif node._op == \"add\":\n",
    "                left, right = node._parents\n",
    "                left._accumulate_grad(node.grad)\n",
    "                right._accumulate_grad(node.grad)\n",
    "\n",
    "            elif node._op == \"sub\":\n",
    "                left, right = node._parents\n",
    "                left._accumulate_grad(node.grad)\n",
    "                right._accumulate_grad(-node.grad)\n",
    "\n",
    "            elif node._op == \"div\":\n",
    "                left, right = node._parents\n",
    "                left._accumulate_grad(node.grad / right.data)\n",
    "                right._accumulate_grad(\n",
    "                    -node.grad * left.data / right.data**2\n",
    "                )\n",
    "            \n",
    "    def __repr__(self):\n",
    "        return f\"Tensor(data={self.data}, grad={self.grad})\"\n",
    "\n",
    "    def __add__(self, other):\n",
    "        out = Tensor(self.data + other.data, requires_grad=self.requires_grad or other.requires_grad)\n",
    "        out._op = \"add\"\n",
    "        out._parents = (self, other)\n",
    "        return out\n",
    "\n",
    "    def __sub__(self, other):\n",
    "        out = Tensor(self.data - other.data, requires_grad=self.requires_grad or other.requires_grad)\n",
    "        out._op = \"sub\"\n",
    "        out._parents = (self, other)\n",
    "        return out\n",
    "    \n",
    "    def __neg__(self):\n",
    "        out = Tensor(-self.data, requires_grad=self.requires_grad)\n",
    "        out._op = \"neg\"\n",
    "        out._parents = (self,)\n",
    "        return out\n",
    "\n",
    "    def __mul__(self, other):\n",
    "        out = Tensor(self.data * other.data, requires_grad=self.requires_grad or other.requires_grad)\n",
    "        out._op = \"mul\"\n",
    "        out._parents = (self, other)    \n",
    "        return out\n",
    "    \n",
    "    def __truediv__(self, other):\n",
    "        out = Tensor(self.data / other.data, requires_grad=self.requires_grad or other.requires_grad)\n",
    "        out._op = \"div\"\n",
    "        out._parents = (self, other)\n",
    "        return out\n",
    "def topo_sort(tensor):\n",
    "    visited = set()\n",
    "    order = []\n",
    "    \n",
    "    def visit(node):\n",
    "        if node in visited:\n",
    "            return\n",
    "\n",
    "        visited.add(node)\n",
    "\n",
    "        for parent in node._parents:\n",
    "            visit(parent)\n",
    "\n",
    "        order.append(node)\n",
    "\n",
    "    visit(tensor)\n",
    "    return order\n",
    "\n",
    "    \n",
    "\n",
    "\n",
    "a = Tensor([2.0], requires_grad=True)\n",
    "b = Tensor([3.0], requires_grad=True)\n",
    "\n",
    "a, b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "44bee163",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "'topo_sort' in globals()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3fa4191f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1.]\n",
      "[3.]\n",
      "[2.]\n"
     ]
    }
   ],
   "source": [
    "c = a * b\n",
    "c.backward()\n",
    "\n",
    "print(c.grad)  # [1.]\n",
    "print(a.grad)  # [3.]\n",
    "print(b.grad)  # [2.]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "2cb4e044",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[-4.]\n",
      "[-2.]\n"
     ]
    }
   ],
   "source": [
    "a = Tensor([2.0], requires_grad=True)\n",
    "b = Tensor([3.0], requires_grad=True)\n",
    "\n",
    "c = a * b\n",
    "d = c + a\n",
    "e = -d\n",
    "\n",
    "e.backward()\n",
    "\n",
    "print(a.grad)  # [-4.]\n",
    "print(b.grad)  # [-2.]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "bfab7ec4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "f = [4.]\n",
      "df/da = [3.33333333]\n",
      "df/db = [-0.66666667]\n"
     ]
    }
   ],
   "source": [
    "a = Tensor([2.0], requires_grad=True)\n",
    "b = Tensor([5.0], requires_grad=True)\n",
    "\n",
    "f = (a * b + a) / (b - a)\n",
    "\n",
    "f.backward()\n",
    "\n",
    "print(\"f =\", f.data)\n",
    "print(\"df/da =\", a.grad)\n",
    "print(\"df/db =\", b.grad)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "965b27a4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "finite difference df/da = 3.3333333333551702\n",
      "finite difference df/db = -0.6666666666710341\n"
     ]
    }
   ],
   "source": [
    "def f_plain(a, b):\n",
    "    return (a * b + a) / (b - a)\n",
    "\n",
    "\n",
    "eps = 1e-5\n",
    "a0 = 2.0\n",
    "b0 = 5.0\n",
    "\n",
    "df_da = (\n",
    "    f_plain(a0 + eps, b0)\n",
    "    - f_plain(a0 - eps, b0)\n",
    ") / (2 * eps)\n",
    "\n",
    "df_db = (\n",
    "    f_plain(a0, b0 + eps)\n",
    "    - f_plain(a0, b0 - eps)\n",
    ") / (2 * eps)\n",
    "\n",
    "print(\"finite difference df/da =\", df_da)\n",
    "print(\"finite difference df/db =\", df_db)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "3a652214",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "autograd df/da = [3.33333333]\n",
      "autograd df/db = [-0.66666667]\n"
     ]
    }
   ],
   "source": [
    "a = Tensor([2.0], requires_grad=True)\n",
    "b = Tensor([5.0], requires_grad=True)\n",
    "\n",
    "f = (a * b + a) / (b - a)\n",
    "\n",
    "f.backward()\n",
    "\n",
    "print(\"autograd df/da =\", a.grad)\n",
    "print(\"autograd df/db =\", b.grad)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "lmlab (3.12.13.final.0)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}