{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# For tips on running notebooks in Google Colab, see\n# https://codelin.vip/beginner/colab\n%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "PyTorch: Defining New autograd Functions\n========================================\n\nA third order polynomial, trained to predict $y=\\sin(x)$ from $-\\pi$ to\n$\\pi$ by minimizing squared Euclidean distance. Instead of writing the\npolynomial as $y=a+bx+cx^2+dx^3$, we write the polynomial as\n$y=a+b P_3(c+dx)$ where $P_3(x)=\\frac{1}{2}\\left(5x^3-3x\\right)$ is the\n[Legendre\npolynomial](https://en.wikipedia.org/wiki/Legendre_polynomials) of\ndegree three.\n\nThis implementation computes the forward pass using operations on\nPyTorch Tensors, and uses PyTorch autograd to compute gradients.\n\nIn this implementation we implement our own custom autograd function to\nperform $P_3'(x)$. By mathematics,\n$P_3'(x)=\\frac{3}{2}\\left(5x^2-1\\right)$\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import torch\nimport math\n\n\nclass LegendrePolynomial3(torch.autograd.Function):\n    \"\"\"\n    We can implement our own custom autograd Functions by subclassing\n    torch.autograd.Function and implementing the forward and backward passes\n    which operate on Tensors.\n    \"\"\"\n\n    @staticmethod\n    def forward(ctx, input):\n        \"\"\"\n        In the forward pass we receive a Tensor containing the input and return\n        a Tensor containing the output. ctx is a context object that can be used\n        to stash information for backward computation. You can cache tensors for\n        use in the backward pass using the ``ctx.save_for_backward`` method. Other\n        objects can be stored directly as attributes on the ctx object, such as\n        ``ctx.my_object = my_object``. Check out `Extending torch.autograd <https://docs.pytorch.org/docs/stable/notes/extending.html#extending-torch-autograd>`_\n        for further details.\n        \"\"\"\n        ctx.save_for_backward(input)\n        return 0.5 * (5 * input ** 3 - 3 * input)\n\n    @staticmethod\n    def backward(ctx, grad_output):\n        \"\"\"\n        In the backward pass we receive a Tensor containing the gradient of the loss\n        with respect to the output, and we need to compute the gradient of the loss\n        with respect to the input.\n        \"\"\"\n        input, = ctx.saved_tensors\n        return grad_output * 1.5 * (5 * input ** 2 - 1)\n\n\ndtype = torch.float\ndevice = torch.device(\"cpu\")\n# device = torch.device(\"cuda:0\")  # Uncomment this to run on GPU\n\n# Create Tensors to hold input and outputs.\n# By default, requires_grad=False, which indicates that we do not need to\n# compute gradients with respect to these Tensors during the backward pass.\nx = torch.linspace(-math.pi, math.pi, 2000, device=device, dtype=dtype)\ny = torch.sin(x)\n\n# Create random Tensors for weights. For this example, we need\n# 4 weights: y = a + b * P3(c + d * x), these weights need to be initialized\n# not too far from the correct result to ensure convergence.\n# Setting requires_grad=True indicates that we want to compute gradients with\n# respect to these Tensors during the backward pass.\na = torch.full((), 0.0, device=device, dtype=dtype, requires_grad=True)\nb = torch.full((), -1.0, device=device, dtype=dtype, requires_grad=True)\nc = torch.full((), 0.0, device=device, dtype=dtype, requires_grad=True)\nd = torch.full((), 0.3, device=device, dtype=dtype, requires_grad=True)\n\nlearning_rate = 5e-6\nfor t in range(2000):\n    # To apply our Function, we use Function.apply method. We alias this as 'P3'.\n    P3 = LegendrePolynomial3.apply\n\n    # Forward pass: compute predicted y using operations; we compute\n    # P3 using our custom autograd operation.\n    y_pred = a + b * P3(c + d * x)\n\n    # Compute and print loss\n    loss = (y_pred - y).pow(2).sum()\n    if t % 100 == 99:\n        print(t, loss.item())\n\n    # Use autograd to compute the backward pass.\n    loss.backward()\n\n    # Update weights using gradient descent\n    with torch.no_grad():\n        a -= learning_rate * a.grad\n        b -= learning_rate * b.grad\n        c -= learning_rate * c.grad\n        d -= learning_rate * d.grad\n\n        # Manually zero the gradients after updating weights\n        a.grad = None\n        b.grad = None\n        c.grad = None\n        d.grad = None\n\nprint(f'Result: y = {a.item()} + {b.item()} * P3({c.item()} + {d.item()} x)')"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "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.12"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}