{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "departmental-hardware",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.autograd.functional import jacobian"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "differential-shock",
   "metadata": {},
   "outputs": [],
   "source": [
    "a = torch.tensor([1,2,3,4.2],requires_grad=False)\n",
    "b = torch.tensor([2,2,2,2.0],requires_grad=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "separated-pursuit",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(x,y):\n",
    "    return (x@y)**2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "french-trunk",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(416.1600, grad_fn=<PowBackward0>)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test(a,b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "adverse-ceremony",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((tensor([81.6000, 81.6000, 81.6000, 81.6000]),\n",
       "  tensor([ 40.8000,  81.6000, 122.4000, 171.3600])),\n",
       " tensor([2., 2., 2., 2.], requires_grad=True),\n",
       " tensor(416.1600, grad_fn=<PowBackward0>))"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "j = jacobian(test,(a,b))\n",
    "j,b,test(a,b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "lovely-apple",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((tensor([-12.8304,  -3.9878,   4.8547,  15.4658]),\n",
       "  tensor([-10.8365, -21.6729, -32.5094, -45.5132])),\n",
       " tensor([ 1.1840,  0.3680, -0.4480, -1.4272], grad_fn=<SubBackward0>),\n",
       " tensor(29.3573, grad_fn=<PowBackward0>))"
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b2 = b - j[1]*b*0.01\n",
    "j2 = jacobian(test,(a,b2))\n",
    "j2,b2,test(a,b2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "stretch-selection",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((tensor([-13.6581,  -4.2906,   5.2787,  17.0284]),\n",
       "  tensor([-11.4119, -22.8239, -34.2358, -47.9301])),\n",
       " tensor([ 1.1968,  0.3760, -0.4626, -1.4922], grad_fn=<SubBackward0>),\n",
       " tensor(32.5580, grad_fn=<PowBackward0>))"
      ]
     },
     "execution_count": 63,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b3 = b2 - j2[1]*b2*0.001\n",
    "j3 = jacobian(test,(a,b3))\n",
    "j3,b3,test(a,b3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "colored-visit",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((tensor([-14.5816,  -4.6324,   5.7628,  18.8361]),\n",
       "  tensor([-12.0461, -24.0921, -36.1382, -50.5935])),\n",
       " tensor([ 1.2105,  0.3846, -0.4784, -1.5637], grad_fn=<SubBackward0>),\n",
       " tensor(36.2769, grad_fn=<PowBackward0>))"
      ]
     },
     "execution_count": 64,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b4 = b3 - j3[1]*b3*0.001\n",
    "j4 = jacobian(test,(a,b4))\n",
    "j4,b4,test(a,b4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "familiar-pizza",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((tensor([-15.6173,  -5.0205,   6.3191,  20.9424]),\n",
       "  tensor([-12.7481, -25.4962, -38.2443, -53.5421])),\n",
       " tensor([ 1.2251,  0.3938, -0.4957, -1.6428], grad_fn=<SubBackward0>),\n",
       " tensor(40.6286, grad_fn=<PowBackward0>))"
      ]
     },
     "execution_count": 65,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b5 = b4 - j4[1]*b4*0.001\n",
    "j5 = jacobian(test,(a,b5))\n",
    "j5,b5,test(a,b5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "brilliant-squad",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((tensor([-15.6173,  -5.0205,   6.3191,  20.9424]),\n",
       "  tensor([-12.7481, -25.4962, -38.2443, -53.5421])),\n",
       " tensor([ 1.2407,  0.4039, -0.5146, -1.7307], grad_fn=<SubBackward0>),\n",
       " tensor(45.7605, grad_fn=<PowBackward0>))"
      ]
     },
     "execution_count": 66,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b6 = b5 - j5[1]*b5*0.001\n",
    "j6 = jacobian(test,(a,b5))\n",
    "j6,b6,test(a,b6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "discrete-engineer",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}