Pytorch 2nd Derivatives.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6f2794b6",
   "metadata": {},
   "source": [
    "# Pytorch 2nd order Derivative computations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 274,
   "id": "09d89c72",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "n_walkers = 10\n",
    "n_particles = 2\n",
    "n_dim = 3\n",
    "torch.set_default_tensor_type(torch.DoubleTensor)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf75a182",
   "metadata": {},
   "source": [
    "Define an input tensor:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 275,
   "id": "653bdee5",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.rand((n_walkers, n_particles, n_dim), requires_grad=True, dtype=torch.float64)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 276,
   "id": "00ff2cff",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([10, 2, 3])"
      ]
     },
     "execution_count": 276,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x.size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 277,
   "id": "113770cf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[0.0466, 0.8930, 0.4716],\n",
      "         [0.1755, 0.4794, 0.7658]],\n",
      "\n",
      "        [[0.4399, 0.7239, 0.3771],\n",
      "         [0.6025, 0.7637, 0.4979]],\n",
      "\n",
      "        [[0.4121, 0.7012, 0.0066],\n",
      "         [0.4676, 0.7512, 0.0072]],\n",
      "\n",
      "        [[0.7359, 0.6154, 0.6321],\n",
      "         [0.5837, 0.6246, 0.3641]],\n",
      "\n",
      "        [[0.4467, 0.0756, 0.1534],\n",
      "         [0.5131, 0.7008, 0.3360]],\n",
      "\n",
      "        [[0.4468, 0.6092, 0.0456],\n",
      "         [0.8910, 0.8574, 0.2396]],\n",
      "\n",
      "        [[0.5234, 0.1006, 0.2207],\n",
      "         [0.8233, 0.5003, 0.9378]],\n",
      "\n",
      "        [[0.2860, 0.5691, 0.8432],\n",
      "         [0.1798, 0.7486, 0.2296]],\n",
      "\n",
      "        [[0.0188, 0.9081, 0.5145],\n",
      "         [0.0465, 0.7274, 0.5099]],\n",
      "\n",
      "        [[0.9027, 0.1038, 0.5196],\n",
      "         [0.4229, 0.5855, 0.0591]]], requires_grad=True)\n"
     ]
    }
   ],
   "source": [
    "print(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc74d23f",
   "metadata": {},
   "source": [
    "Define a model to differentiate:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 278,
   "id": "80bf6504",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DeepSets(torch.nn.Module):\n",
    "    \n",
    "    def __init__(self):\n",
    "        torch.nn.Module.__init__(self)\n",
    "        \n",
    "        self.individual_net = torch.nn.ModuleList([\n",
    "            torch.nn.Linear(n_dim, 16, bias=False),\n",
    "            torch.nn.Linear(16, 32, bias=False)\n",
    "        ]\n",
    "        )\n",
    "        \n",
    "        self.aggregate_net = torch.nn.ModuleList([\n",
    "            torch.nn.Linear(32, 16, bias=False),\n",
    "            torch.nn.Linear(16, 1, bias=False)\n",
    "        ]\n",
    "        )\n",
    "        \n",
    "    def forward(self, inputs):\n",
    "        \n",
    "        #split the inputs along the particle dimension (1):\n",
    "        particles = torch.chunk(inputs, n_particles, axis=1)\n",
    "        \n",
    "        particles = [torch.reshape(p, (-1,n_dim)) for p in particles]\n",
    "        \n",
    "        individuals = []\n",
    "        for p in particles:\n",
    "            this_i = p\n",
    "            for l in self.individual_net:\n",
    "                this_i = torch.tanh(l(this_i))\n",
    "            individuals.append(this_i)\n",
    "#         individuals = [self.individual_net(p) for p in particles]\n",
    "\n",
    "        concatd = torch.stack(individuals, dim=-1)\n",
    "        \n",
    "        \n",
    "        # Sum over the latent space:\n",
    "        summed = torch.sum(concatd, dim=-1)\n",
    "        \n",
    "        output = summed\n",
    "        for l in self.aggregate_net:\n",
    "            output = l(output)\n",
    "#         output = self.(summed)\n",
    "        \n",
    "        return output.reshape((-1))\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 279,
   "id": "3ba15e09",
   "metadata": {},
   "outputs": [],
   "source": [
    "d = DeepSets()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 280,
   "id": "f716232c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "print(x.requires_grad)\n",
    "o = d(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 281,
   "id": "9efa9a53",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([10])"
      ]
     },
     "execution_count": 281,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "o.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 282,
   "id": "b8f968b9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<ReshapeAliasBackward0 at 0x7fcbd1d5ec88>"
      ]
     },
     "execution_count": 282,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "o.grad_fn"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8819b7d",
   "metadata": {},
   "source": [
    "# How to compute derivatives"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 283,
   "id": "b50f8dcb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# First, get the first derivative at every particle coordinate:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 284,
   "id": "559185c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "w_of_x, jvp = torch.autograd.functional.vjp(d.forward, x, torch.ones((n_walkers)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 285,
   "id": "8d1cb69c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([10, 2, 3])"
      ]
     },
     "execution_count": 285,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jvp.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 286,
   "id": "9ad5e55b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], grad_fn=<SubBackward0>)"
      ]
     },
     "execution_count": 286,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "o - w_of_x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04f21811",
   "metadata": {},
   "source": [
    "Check the derivatives with numerical approximation (NOTE: it helps to use double precision for this):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 287,
   "id": "78bc0dd9",
   "metadata": {},
   "outputs": [],
   "source": [
    "kick = torch.zeros(x.shape)\n",
    "dim = torch.randint(low=0,high=n_dim, size=()).item()\n",
    "particle = torch.randint(low=0,high=n_particles, size=()).item()\n",
    "kick_size = 1e-5\n",
    "kick[:,particle, dim] = kick_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 288,
   "id": "e2482fe3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([10])\n",
      "torch.Size([10])\n"
     ]
    }
   ],
   "source": [
    "numerical_dw_dx = torch.reshape((d(x + kick) - d(x - kick))/(2*kick_size), (-1,))\n",
    "print(numerical_dw_dx.shape)\n",
    "print(jvp[:,particle, dim].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 289,
   "id": "85762524",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0.0550, 0.0598, 0.0646, 0.0570, 0.0648, 0.0648, 0.0643, 0.0523, 0.0540,\n",
      "        0.0607], grad_fn=<ReshapeAliasBackward0>)\n",
      "tensor([0.0550, 0.0598, 0.0646, 0.0570, 0.0648, 0.0648, 0.0643, 0.0523, 0.0540,\n",
      "        0.0607])\n",
      "tensor([ 1.5443e-12,  3.0219e-14,  3.2294e-12,  1.6786e-13,  2.2961e-12,\n",
      "         1.0515e-12, -1.6431e-12, -7.0966e-13, -9.7358e-13, -1.8749e-12],\n",
      "       grad_fn=<SubBackward0>)\n"
     ]
    }
   ],
   "source": [
    "print(numerical_dw_dx)\n",
    "print(jvp[:,particle, dim])\n",
    "print(numerical_dw_dx - jvp[:,particle, dim] )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 290,
   "id": "3d524ebe",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "print(x.requires_grad)\n",
    "o = d(x)\n",
    "dw_dx = torch.autograd.grad(o, x, grad_outputs=torch.ones((n_walkers)), retain_graph=True, create_graph=True)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 291,
   "id": "ca32859b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[-0.0445,  0.0355,  0.0550],\n",
       "         [-0.0443,  0.0454,  0.0541]],\n",
       "\n",
       "        [[-0.0479,  0.0404,  0.0598],\n",
       "         [-0.0482,  0.0376,  0.0582]],\n",
       "\n",
       "        [[-0.0490,  0.0436,  0.0646],\n",
       "         [-0.0495,  0.0417,  0.0646]],\n",
       "\n",
       "        [[-0.0473,  0.0411,  0.0570],\n",
       "         [-0.0478,  0.0430,  0.0611]],\n",
       "\n",
       "        [[-0.0452,  0.0557,  0.0648],\n",
       "         [-0.0482,  0.0411,  0.0609]],\n",
       "\n",
       "        [[-0.0484,  0.0458,  0.0648],\n",
       "         [-0.0489,  0.0343,  0.0622]],\n",
       "\n",
       "        [[-0.0452,  0.0551,  0.0643],\n",
       "         [-0.0462,  0.0427,  0.0523]],\n",
       "\n",
       "        [[-0.0449,  0.0424,  0.0523],\n",
       "         [-0.0470,  0.0418,  0.0607]],\n",
       "\n",
       "        [[-0.0438,  0.0346,  0.0540],\n",
       "         [-0.0440,  0.0404,  0.0559]],\n",
       "\n",
       "        [[-0.0442,  0.0527,  0.0607],\n",
       "         [-0.0481,  0.0465,  0.0647]]], grad_fn=<CatBackward0>)"
      ]
     },
     "execution_count": 291,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dw_dx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 292,
   "id": "8b705785",
   "metadata": {},
   "outputs": [],
   "source": [
    "d2w_dx2 = torch.autograd.grad(dw_dx, x, grad_outputs=torch.ones_like(x))[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 293,
   "id": "5a716b57",
   "metadata": {},
   "outputs": [],
   "source": [
    "# How about a second order derivative?  Recompute the first order and create a graph:\n",
    "# l = lambda x : d.forward(x).sum()\n",
    "# d2w_dx2, hvp = torch.autograd.functional.hvp(l, x, torch.ones_like(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 294,
   "id": "010aacb5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[-0.0091, -0.0434, -0.0178],\n",
      "         [-0.0039, -0.0328, -0.0233]],\n",
      "\n",
      "        [[-0.0053, -0.0441, -0.0187],\n",
      "         [-0.0038, -0.0449, -0.0197]],\n",
      "\n",
      "        [[-0.0055, -0.0408, -0.0111],\n",
      "         [-0.0051, -0.0422, -0.0118]],\n",
      "\n",
      "        [[-0.0022, -0.0417, -0.0208],\n",
      "         [-0.0040, -0.0418, -0.0182]],\n",
      "\n",
      "        [[-0.0039, -0.0115, -0.0055],\n",
      "         [-0.0047, -0.0436, -0.0181]],\n",
      "\n",
      "        [[-0.0054, -0.0382, -0.0108],\n",
      "         [ 0.0004, -0.0447, -0.0175]],\n",
      "\n",
      "        [[-0.0041, -0.0147, -0.0081],\n",
      "         [-0.0007, -0.0373, -0.0219]],\n",
      "\n",
      "        [[-0.0049, -0.0359, -0.0222],\n",
      "         [-0.0069, -0.0432, -0.0164]],\n",
      "\n",
      "        [[-0.0094, -0.0428, -0.0176],\n",
      "         [-0.0072, -0.0404, -0.0195]],\n",
      "\n",
      "        [[-0.0027, -0.0202, -0.0148],\n",
      "         [-0.0055, -0.0375, -0.0109]]])\n"
     ]
    }
   ],
   "source": [
    "print(d2w_dx2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 298,
   "id": "fd73ca86",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([10, 2, 3])"
      ]
     },
     "execution_count": 298,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "d2w_dx2.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 299,
   "id": "1dfdb2f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# d2w_dx2, vhp = torch.autograd.functional.vhp(l, x, torch.ones_like(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 300,
   "id": "4ad86cea",
   "metadata": {},
   "outputs": [],
   "source": [
    "# torch.max(hvp - hvp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 301,
   "id": "9ced991c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# torch.min(hvp - hvp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 302,
   "id": "c1baa946",
   "metadata": {},
   "outputs": [],
   "source": [
    "kick = torch.zeros(x.shape)\n",
    "dim = torch.randint(low=0,high=n_dim, size=()).item()\n",
    "particle = torch.randint(low=0,high=n_particles, size=()).item()\n",
    "kick_size = 1e-4\n",
    "kick[:,particle, dim] = kick_size\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 303,
   "id": "8c851ab7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([-0.0135, -0.0043, -0.0024,  0.0005,  0.0027, -0.0015,  0.0028, -0.0079,\n",
      "        -0.0143,  0.0050], grad_fn=<DivBackward0>)\n",
      "tensor([-0.0135, -0.0043, -0.0024,  0.0005,  0.0027, -0.0015,  0.0028, -0.0079,\n",
      "        -0.0143,  0.0050], grad_fn=<DivBackward0>)\n",
      "tensor([1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000,\n",
      "        1.0000], grad_fn=<DivBackward0>)\n"
     ]
    }
   ],
   "source": [
    "numerical_d2w_dx2_order1 = (d(x + kick) - 2*d(x) + d(x-kick)) / (kick_size * kick_size)\n",
    "numerical_d2w_dx2_order2 = (-d(x-2*kick) + 16*d(x-kick) - 30*d(x) + 16*d(x+kick) - d(x+2*kick)) / (12*kick_size * kick_size)\n",
    "print(numerical_d2w_dx2_order1)\n",
    "print(numerical_d2w_dx2_order2)\n",
    "\n",
    "print(numerical_d2w_dx2_order1 / numerical_d2w_dx2_order2)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 304,
   "id": "c3aa4190",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([10])\n",
      "torch.Size([10])\n",
      "tensor([-0.0091, -0.0053, -0.0055, -0.0022, -0.0039, -0.0054, -0.0041, -0.0049,\n",
      "        -0.0094, -0.0027])\n"
     ]
    }
   ],
   "source": [
    "numerical_d2w_dx2 = numerical_d2w_dx2_order2.reshape(-1)\n",
    "print(numerical_d2w_dx2.shape)\n",
    "print(d2w_dx2[:,particle, dim].shape)\n",
    "print(d2w_dx2[:,particle,dim])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 305,
   "id": "5278bfa0",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([-0.0135, -0.0043, -0.0024,  0.0005,  0.0027, -0.0015,  0.0028, -0.0079,\n",
      "        -0.0143,  0.0050], grad_fn=<ReshapeAliasBackward0>)\n",
      "tensor([-0.0091, -0.0053, -0.0055, -0.0022, -0.0039, -0.0054, -0.0041, -0.0049,\n",
      "        -0.0094, -0.0027])\n",
      "tensor([-0.0044,  0.0010,  0.0032,  0.0026,  0.0066,  0.0039,  0.0069, -0.0031,\n",
      "        -0.0049,  0.0078], grad_fn=<SubBackward0>)\n",
      "tensor([ 1.4898,  0.8075,  0.4285, -0.2112, -0.6737,  0.2779, -0.6931,  1.6363,\n",
      "         1.5185, -1.8507], grad_fn=<DivBackward0>)\n"
     ]
    }
   ],
   "source": [
    "\n",
    "print(numerical_d2w_dx2)\n",
    "print(d2w_dx2[:,particle, dim])\n",
    "print(numerical_d2w_dx2 - d2w_dx2[:,particle, dim] )\n",
    "print(numerical_d2w_dx2 / d2w_dx2[:,particle, dim] )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 306,
   "id": "57ca8023",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "torch.Size([10, 2, 3])\n",
      "tensor([[[-0.0445,  0.0355,  0.0550],\n",
      "         [-0.0443,  0.0454,  0.0541]],\n",
      "\n",
      "        [[-0.0479,  0.0404,  0.0598],\n",
      "         [-0.0482,  0.0376,  0.0582]],\n",
      "\n",
      "        [[-0.0490,  0.0436,  0.0646],\n",
      "         [-0.0495,  0.0417,  0.0646]],\n",
      "\n",
      "        [[-0.0473,  0.0411,  0.0570],\n",
      "         [-0.0478,  0.0430,  0.0611]],\n",
      "\n",
      "        [[-0.0452,  0.0557,  0.0648],\n",
      "         [-0.0482,  0.0411,  0.0609]],\n",
      "\n",
      "        [[-0.0484,  0.0458,  0.0648],\n",
      "         [-0.0489,  0.0343,  0.0622]],\n",
      "\n",
      "        [[-0.0452,  0.0551,  0.0643],\n",
      "         [-0.0462,  0.0427,  0.0523]],\n",
      "\n",
      "        [[-0.0449,  0.0424,  0.0523],\n",
      "         [-0.0470,  0.0418,  0.0607]],\n",
      "\n",
      "        [[-0.0438,  0.0346,  0.0540],\n",
      "         [-0.0440,  0.0404,  0.0559]],\n",
      "\n",
      "        [[-0.0442,  0.0527,  0.0607],\n",
      "         [-0.0481,  0.0465,  0.0647]]], grad_fn=<CatBackward0>)\n"
     ]
    }
   ],
   "source": [
    "print(x.requires_grad)\n",
    "o = d(x)\n",
    "dw_dx = torch.autograd.grad(o,x,grad_outputs=torch.ones((n_walkers)), retain_graph=True, create_graph=True)[0]\n",
    "print(dw_dx.shape)\n",
    "print(dw_dx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 307,
   "id": "589ec5c7",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([10, 2, 3])\n"
     ]
    }
   ],
   "source": [
    "d2w_dx2_slow = torch.zeros_like(x)\n",
    "print(x.shape)\n",
    "for i_part in range(n_particles):\n",
    "    for i_dim in range(n_dim):\n",
    "        dw_dx_ij = dw_dx[:,i_part,i_dim]\n",
    "\n",
    "        d2w_dx2_ij = torch.autograd.grad(dw_dx_ij, x, grad_outputs=torch.ones([n_walkers,]), retain_graph=True)[0]\n",
    "        \n",
    "        d2w_dx2_slow[:,i_part, i_dim] = d2w_dx2_ij[:,i_part,i_dim]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 308,
   "id": "b3744479",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([10, 2, 3])"
      ]
     },
     "execution_count": 308,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "d2w_dx2_slow.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 309,
   "id": "846a432d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-0.0135, -0.0043, -0.0024,  0.0005,  0.0027, -0.0015,  0.0028, -0.0079,\n",
       "        -0.0143,  0.0050], grad_fn=<DivBackward0>)"
      ]
     },
     "execution_count": 309,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "numerical_d2w_dx2_order2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 310,
   "id": "d71048e1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-0.0135, -0.0043, -0.0024,  0.0005,  0.0027, -0.0015,  0.0028, -0.0079,\n",
       "        -0.0143,  0.0050])"
      ]
     },
     "execution_count": 310,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "d2w_dx2_slow[:,particle,dim]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 318,
   "id": "fc7a8268",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "d2w_dx2 = torch.autograd.grad(dw_dx, x, grad_outputs=torch.ones((n_walkers, n_particles, n_dim)), retain_graph=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 313,
   "id": "10d47383",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[-0.0135, -0.0326, -0.0155],\n",
       "         [-0.0084, -0.0257, -0.0183]],\n",
       "\n",
       "        [[-0.0043, -0.0327, -0.0155],\n",
       "         [-0.0019, -0.0339, -0.0165]],\n",
       "\n",
       "        [[-0.0024, -0.0292, -0.0095],\n",
       "         [-0.0016, -0.0300, -0.0099]],\n",
       "\n",
       "        [[ 0.0005, -0.0319, -0.0175],\n",
       "         [-0.0011, -0.0311, -0.0150]],\n",
       "\n",
       "        [[ 0.0027, -0.0074, -0.0048],\n",
       "         [-0.0026, -0.0322, -0.0150]],\n",
       "\n",
       "        [[-0.0015, -0.0277, -0.0095],\n",
       "         [ 0.0048, -0.0332, -0.0136]],\n",
       "\n",
       "        [[ 0.0028, -0.0096, -0.0068],\n",
       "         [ 0.0013, -0.0286, -0.0184]],\n",
       "\n",
       "        [[-0.0079, -0.0285, -0.0184],\n",
       "         [-0.0085, -0.0314, -0.0136]],\n",
       "\n",
       "        [[-0.0143, -0.0324, -0.0155],\n",
       "         [-0.0121, -0.0310, -0.0163]],\n",
       "\n",
       "        [[ 0.0050, -0.0129, -0.0119],\n",
       "         [-0.0018, -0.0272, -0.0096]]])"
      ]
     },
     "execution_count": 313,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "d2w_dx2_slow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 319,
   "id": "5c127ff2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[[-0.0091, -0.0434, -0.0178],\n",
       "          [-0.0039, -0.0328, -0.0233]],\n",
       " \n",
       "         [[-0.0053, -0.0441, -0.0187],\n",
       "          [-0.0038, -0.0449, -0.0197]],\n",
       " \n",
       "         [[-0.0055, -0.0408, -0.0111],\n",
       "          [-0.0051, -0.0422, -0.0118]],\n",
       " \n",
       "         [[-0.0022, -0.0417, -0.0208],\n",
       "          [-0.0040, -0.0418, -0.0182]],\n",
       " \n",
       "         [[-0.0039, -0.0115, -0.0055],\n",
       "          [-0.0047, -0.0436, -0.0181]],\n",
       " \n",
       "         [[-0.0054, -0.0382, -0.0108],\n",
       "          [ 0.0004, -0.0447, -0.0175]],\n",
       " \n",
       "         [[-0.0041, -0.0147, -0.0081],\n",
       "          [-0.0007, -0.0373, -0.0219]],\n",
       " \n",
       "         [[-0.0049, -0.0359, -0.0222],\n",
       "          [-0.0069, -0.0432, -0.0164]],\n",
       " \n",
       "         [[-0.0094, -0.0428, -0.0176],\n",
       "          [-0.0072, -0.0404, -0.0195]],\n",
       " \n",
       "         [[-0.0027, -0.0202, -0.0148],\n",
       "          [-0.0055, -0.0375, -0.0109]]]),)"
      ]
     },
     "execution_count": 319,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "d2w_dx2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 330,
   "id": "82bc17f8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([10, 2, 3])\n",
      "torch.Size([10])\n",
      "torch.Size([10, 2, 3])\n"
     ]
    },
    {
     "ename": "RuntimeError",
     "evalue": "Mismatch in shape: grad_output[0] has a shape of torch.Size([10, 2, 3]) and output[0] has a shape of torch.Size([10]).",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[0;32m/var/folders/f_/5dsy8_2x1wd93w36cdk7twz9s8nq7s/T/ipykernel_18479/4223813038.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdw_dx_ij_test\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 5\u001b[0;31m \u001b[0md2w_dx2_test\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mautograd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgrad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdw_dx_ij_test\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgrad_outputs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mones\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mn_walkers\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m3\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m~/SOFTWARE/larcv-env/lib/python3.7/site-packages/torch/autograd/__init__.py\u001b[0m in \u001b[0;36mgrad\u001b[0;34m(outputs, inputs, grad_outputs, retain_graph, create_graph, only_inputs, allow_unused, is_grads_batched)\u001b[0m\n\u001b[1;32m    258\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    259\u001b[0m     \u001b[0mgrad_outputs_\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_tensor_or_tensors_to_tuple\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mgrad_outputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moutputs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 260\u001b[0;31m     \u001b[0mgrad_outputs_\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_make_grads\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moutputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgrad_outputs_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mis_grads_batched\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mis_grads_batched\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    261\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    262\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mretain_graph\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/SOFTWARE/larcv-env/lib/python3.7/site-packages/torch/autograd/__init__.py\u001b[0m in \u001b[0;36m_make_grads\u001b[0;34m(outputs, grads, is_grads_batched)\u001b[0m\n\u001b[1;32m     52\u001b[0m                                        \u001b[0;34m+\u001b[0m \u001b[0mstr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mgrad\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;34m\" and output[\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     53\u001b[0m                                        \u001b[0;34m+\u001b[0m \u001b[0mstr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moutputs\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mindex\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mout\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;34m\"] has a shape of \"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 54\u001b[0;31m                                        + str(out.shape) + \".\")\n\u001b[0m\u001b[1;32m     55\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mout\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdtype\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mis_complex\u001b[0m \u001b[0;34m!=\u001b[0m \u001b[0mgrad\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdtype\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mis_complex\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     56\u001b[0m                 raise RuntimeError(\"For complex Tensors, both grad_output and output\"\n",
      "\u001b[0;31mRuntimeError\u001b[0m: Mismatch in shape: grad_output[0] has a shape of torch.Size([10, 2, 3]) and output[0] has a shape of torch.Size([10])."
     ]
    }
   ],
   "source": [
    "dw_dx_ij_test = dw_dx[:,0,0]\n",
    "print(dw_dx.shape)\n",
    "print(dw_dx_ij_test.shape)\n",
    "print(x.shape)\n",
    "d2w_dx2_test = torch.autograd.grad(dw_dx_ij_test, x, grad_outputs=torch.ones([n_walkers,2,3]), retain_graph=True)[0]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 347,
   "id": "8d9c332d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([10, 2, 3])\n",
      "torch.Size([10, 2, 3])\n"
     ]
    }
   ],
   "source": [
    "print(dw_dx.shape)\n",
    "print(x.shape)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "393065f2",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 349,
   "id": "adb59bb4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2"
      ]
     },
     "execution_count": 349,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dw_dx.shape[-2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 423,
   "id": "ab698b17",
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = x\n",
    "o = d(inputs)\n",
    "dw_dx = torch.autograd.grad(o,inputs,grad_outputs=torch.ones([n_walkers]), retain_graph=True, create_graph=True)[0]\n",
    "summed = torch.sum(dw_dx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 425,
   "id": "0d5ac1cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "f = summed.backward()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 426,
   "id": "139eb2c0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "None\n"
     ]
    }
   ],
   "source": [
    "print(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 427,
   "id": "a157ef86",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[0.6712, 1.3317, 1.1502],\n",
      "         [0.4664, 1.2762, 1.2711]]])\n",
      "tensor([[[0.0466, 0.8930, 0.4716],\n",
      "         [0.1755, 0.4794, 0.7658]],\n",
      "\n",
      "        [[0.4399, 0.7239, 0.3771],\n",
      "         [0.6025, 0.7637, 0.4979]],\n",
      "\n",
      "        [[0.4121, 0.7012, 0.0066],\n",
      "         [0.4676, 0.7512, 0.0072]],\n",
      "\n",
      "        [[0.7359, 0.6154, 0.6321],\n",
      "         [0.5837, 0.6246, 0.3641]],\n",
      "\n",
      "        [[0.4467, 0.0756, 0.1534],\n",
      "         [0.5131, 0.7008, 0.3360]],\n",
      "\n",
      "        [[0.4468, 0.6092, 0.0456],\n",
      "         [0.8910, 0.8574, 0.2396]],\n",
      "\n",
      "        [[0.5234, 0.1006, 0.2207],\n",
      "         [0.8233, 0.5003, 0.9378]],\n",
      "\n",
      "        [[0.2860, 0.5691, 0.8432],\n",
      "         [0.1798, 0.7486, 0.2296]],\n",
      "\n",
      "        [[0.0188, 0.9081, 0.5145],\n",
      "         [0.0465, 0.7274, 0.5099]],\n",
      "\n",
      "        [[0.9027, 0.1038, 0.5196],\n",
      "         [0.4229, 0.5855, 0.0591]]], requires_grad=True)\n"
     ]
    }
   ],
   "source": [
    "print(d2w_dx2_ein / d2w_dx2_slow[0])\n",
    "print(inputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 362,
   "id": "dd096c04",
   "metadata": {},
   "outputs": [],
   "source": [
    "f = lambda x : d(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 370,
   "id": "1a357b66",
   "metadata": {},
   "outputs": [],
   "source": [
    "_, g = torch.autograd.functional.jvp(f, x, v=torch.ones_like(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 372,
   "id": "82e3cf06",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([10])"
      ]
     },
     "execution_count": 372,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "g.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 366,
   "id": "6d96d0ea",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[-0.0445,  0.0355,  0.0550],\n",
       "         [-0.0443,  0.0454,  0.0541]],\n",
       "\n",
       "        [[-0.0479,  0.0404,  0.0598],\n",
       "         [-0.0482,  0.0376,  0.0582]],\n",
       "\n",
       "        [[-0.0490,  0.0436,  0.0646],\n",
       "         [-0.0495,  0.0417,  0.0646]],\n",
       "\n",
       "        [[-0.0473,  0.0411,  0.0570],\n",
       "         [-0.0478,  0.0430,  0.0611]],\n",
       "\n",
       "        [[-0.0452,  0.0557,  0.0648],\n",
       "         [-0.0482,  0.0411,  0.0609]],\n",
       "\n",
       "        [[-0.0484,  0.0458,  0.0648],\n",
       "         [-0.0489,  0.0343,  0.0622]],\n",
       "\n",
       "        [[-0.0452,  0.0551,  0.0643],\n",
       "         [-0.0462,  0.0427,  0.0523]],\n",
       "\n",
       "        [[-0.0449,  0.0424,  0.0523],\n",
       "         [-0.0470,  0.0418,  0.0607]],\n",
       "\n",
       "        [[-0.0438,  0.0346,  0.0540],\n",
       "         [-0.0440,  0.0404,  0.0559]],\n",
       "\n",
       "        [[-0.0442,  0.0527,  0.0607],\n",
       "         [-0.0481,  0.0465,  0.0647]]], grad_fn=<CatBackward0>)"
      ]
     },
     "execution_count": 366,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dw_dx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 408,
   "id": "9e1816e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "func = lambda x : d(x)\n",
    "jacf = lambda x : torch.autograd.functional.jacobian(func, x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 409,
   "id": "cae6c86e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[-0.0445,  0.0355,  0.0550],\n",
       "          [-0.0443,  0.0454,  0.0541]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]]],\n",
       "\n",
       "\n",
       "        [[[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[-0.0479,  0.0404,  0.0598],\n",
       "          [-0.0482,  0.0376,  0.0582]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]]],\n",
       "\n",
       "\n",
       "        [[[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[-0.0490,  0.0436,  0.0646],\n",
       "          [-0.0495,  0.0417,  0.0646]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]]],\n",
       "\n",
       "\n",
       "        [[[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[-0.0473,  0.0411,  0.0570],\n",
       "          [-0.0478,  0.0430,  0.0611]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]]],\n",
       "\n",
       "\n",
       "        [[[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[-0.0452,  0.0557,  0.0648],\n",
       "          [-0.0482,  0.0411,  0.0609]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]]],\n",
       "\n",
       "\n",
       "        [[[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[-0.0484,  0.0458,  0.0648],\n",
       "          [-0.0489,  0.0343,  0.0622]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]]],\n",
       "\n",
       "\n",
       "        [[[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[-0.0452,  0.0551,  0.0643],\n",
       "          [-0.0462,  0.0427,  0.0523]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]]],\n",
       "\n",
       "\n",
       "        [[[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[-0.0449,  0.0424,  0.0523],\n",
       "          [-0.0470,  0.0418,  0.0607]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]]],\n",
       "\n",
       "\n",
       "        [[[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[-0.0438,  0.0346,  0.0540],\n",
       "          [-0.0440,  0.0404,  0.0559]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]]],\n",
       "\n",
       "\n",
       "        [[[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[ 0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000]],\n",
       "\n",
       "         [[-0.0442,  0.0527,  0.0607],\n",
       "          [-0.0481,  0.0465,  0.0647]]]])"
      ]
     },
     "execution_count": 409,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jacf(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6a34b2c",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}