Pytorch Jacobian.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "66e6ac11",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch, numpy\n",
    "n_walkers = 20\n",
    "n_particles = 2\n",
    "n_dim = 3\n",
    "torch.set_default_tensor_type(torch.DoubleTensor)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bdba9406",
   "metadata": {},
   "source": [
    "Input tensor:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 136,
   "id": "e6f8990a",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.rand((n_walkers, n_particles, n_dim), requires_grad=True, dtype=torch.float64)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0edd3288",
   "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": 4,
   "id": "74256f39",
   "metadata": {},
   "outputs": [],
   "source": [
    "d = DeepSets()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4c496fab",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "print(x.requires_grad)\n",
    "o = d(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "bf293c96",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([20])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "o.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f81c486f",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d5b72e0e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<ReshapeAliasBackward0 at 0x7f8054234080>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "o.grad_fn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "86c07d41",
   "metadata": {},
   "outputs": [],
   "source": [
    "for param_tensor in d.parameters():\n",
    "    jac_i = torch.autograd.functional.jacobian(d, x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "507d56f7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([20, 20, 2, 3])\n"
     ]
    }
   ],
   "source": [
    "print(jac_i.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d33f00b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "o = d(x)\n",
    "params = list(d.parameters())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "0a8cc28b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([-0.0758, -0.0851, -0.0414, -0.0907, -0.0402, -0.0218, -0.0074, -0.0455,\n",
      "         0.0101, -0.0610, -0.0351, -0.0358, -0.0311, -0.0416, -0.0185, -0.0245,\n",
      "        -0.0054, -0.0137,  0.0060, -0.0059], grad_fn=<ReshapeAliasBackward0>)\n"
     ]
    }
   ],
   "source": [
    "print(o)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "4d0b1c02",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(tensor([1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0.]), tensor([0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0.]), tensor([0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0.]), tensor([0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0.]), tensor([0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0.]), tensor([0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0.]), tensor([0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0.]), tensor([0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0.]), tensor([0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0.]), tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0.]), tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0.]), tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0.]), tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0.,\n",
      "        0., 0.]), tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0.,\n",
      "        0., 0.]), tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0.,\n",
      "        0., 0.]), tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0.,\n",
      "        0., 0.]), tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0.,\n",
      "        0., 0.]), tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1.,\n",
      "        0., 0.]), tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        1., 0.]), tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 1.]))\n"
     ]
    }
   ],
   "source": [
    "outputs = torch.tensor_split(torch.flatten(torch.eye(n_walkers)), n_walkers, axis=0)\n",
    "print(outputs)\n",
    "jac = torch.autograd.grad(o, params, retain_graph=True, grad_outputs=outputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "08d38a0e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-8.9194e-04,  7.4605e-04, -3.8275e-03,  4.8174e-03, -2.7491e-04,\n",
       "         -3.0629e-03,  1.8783e-03, -1.6623e-04,  4.2345e-03,  5.2004e-03,\n",
       "          2.9563e-03, -2.3737e-03, -1.1785e-03, -8.6339e-04,  5.1584e-03,\n",
       "         -1.3632e-06],\n",
       "        [ 1.3119e-02, -1.0909e-02,  5.6119e-02, -7.0482e-02,  3.9984e-03,\n",
       "          4.4917e-02, -2.7500e-02,  2.3416e-03, -6.1993e-02, -7.6153e-02,\n",
       "         -4.3278e-02,  3.4787e-02,  1.7267e-02,  1.2622e-02, -7.5508e-02,\n",
       "         -3.0392e-05],\n",
       "        [-9.8766e-03,  8.3813e-03, -4.2718e-02,  5.4045e-02, -3.1284e-03,\n",
       "         -3.4166e-02,  2.1036e-02, -2.0336e-03,  4.7435e-02,  5.8220e-02,\n",
       "          3.3120e-02, -2.6523e-02, -1.3173e-02, -9.7051e-03,  5.7804e-02,\n",
       "         -1.0915e-04],\n",
       "        [-7.2110e-03,  5.9622e-03, -3.0751e-02,  3.8541e-02, -2.1738e-03,\n",
       "         -2.4617e-02,  1.5048e-02, -1.2323e-03,  3.3920e-02,  4.1677e-02,\n",
       "          2.3678e-02, -1.9053e-02, -9.4559e-03, -6.8966e-03,  4.1309e-02,\n",
       "          4.3362e-05],\n",
       "        [ 5.1995e-03, -4.4020e-03,  2.2460e-02, -2.8392e-02,  1.6397e-03,\n",
       "          1.7965e-02, -1.1054e-02,  1.0540e-03, -2.4925e-02, -3.0595e-02,\n",
       "         -1.7403e-02,  1.3942e-02,  6.9244e-03,  5.0968e-03, -3.0372e-02,\n",
       "          4.9357e-05],\n",
       "        [ 1.2302e-03, -1.0238e-03,  5.2647e-03, -6.6138e-03,  3.7549e-04,\n",
       "          4.2136e-03, -2.5803e-03,  2.2083e-04, -5.8168e-03, -7.1452e-03,\n",
       "         -4.0608e-03,  3.2636e-03,  1.6200e-03,  1.1845e-03, -7.0851e-03,\n",
       "         -2.2389e-06],\n",
       "        [ 4.1269e-03, -3.4369e-03,  1.7668e-02, -2.2202e-02,  1.2614e-03,\n",
       "          1.4140e-02, -8.6608e-03,  7.4498e-04, -1.9525e-02, -2.3983e-02,\n",
       "         -1.3630e-02,  1.0953e-02,  5.4372e-03,  3.9767e-03, -2.3782e-02,\n",
       "         -5.4606e-06],\n",
       "        [-8.3727e-03,  6.9560e-03, -3.5798e-02,  4.4945e-02, -2.5474e-03,\n",
       "         -2.8653e-02,  1.7538e-02, -1.4844e-03,  3.9535e-02,  4.8567e-02,\n",
       "          2.7600e-02, -2.2189e-02, -1.1014e-02, -8.0478e-03,  4.8154e-02,\n",
       "          2.4238e-05],\n",
       "        [ 4.1780e-03, -3.5236e-03,  1.8009e-02, -2.2734e-02,  1.3080e-03,\n",
       "          1.4407e-02, -8.8554e-03,  8.2509e-04, -1.9967e-02, -2.4513e-02,\n",
       "         -1.3940e-02,  1.1176e-02,  5.5500e-03,  4.0791e-03, -2.4328e-02,\n",
       "          2.9029e-05],\n",
       "        [ 4.4338e-03, -3.6661e-03,  1.8908e-02, -2.3699e-02,  1.3367e-03,\n",
       "          1.5137e-02, -9.2528e-03,  7.5803e-04, -2.0857e-02, -2.5627e-02,\n",
       "         -1.4560e-02,  1.1716e-02,  5.8143e-03,  4.2407e-03, -2.5401e-02,\n",
       "         -2.6517e-05],\n",
       "        [-1.4931e-03,  1.2392e-03, -6.3805e-03,  8.0079e-03, -4.5342e-04,\n",
       "         -5.1072e-03,  3.1252e-03, -2.6273e-04,  7.0449e-03,  8.6547e-03,\n",
       "          4.9180e-03, -3.9545e-03, -1.9628e-03, -1.4337e-03,  8.5804e-03,\n",
       "          5.2950e-06],\n",
       "        [ 1.3172e-02, -1.0937e-02,  5.6300e-02, -7.0671e-02,  4.0031e-03,\n",
       "          4.5064e-02, -2.7579e-02,  2.3250e-03, -6.2169e-02, -7.6374e-02,\n",
       "         -4.3400e-02,  3.4895e-02,  1.7320e-02,  1.2653e-02, -7.5720e-02,\n",
       "         -4.3190e-05],\n",
       "        [-1.9289e-02,  1.6079e-02, -8.2622e-02,  1.0386e-01, -5.9065e-03,\n",
       "         -6.6123e-02,  4.0511e-02, -3.5063e-03,  9.1326e-02,  1.1217e-01,\n",
       "          6.3757e-02, -5.1226e-02, -2.5429e-02, -1.8605e-02,  1.1124e-01,\n",
       "          1.3721e-05],\n",
       "        [ 1.0717e-02, -9.0070e-03,  4.6110e-02, -5.8133e-02,  3.3332e-03,\n",
       "          3.6892e-02, -2.2653e-02,  2.0659e-03, -5.1075e-02, -6.2713e-02,\n",
       "         -3.5659e-02,  2.8607e-02,  1.4204e-02,  1.0426e-02, -6.2225e-02,\n",
       "          4.9798e-05],\n",
       "        [ 8.9310e-04, -7.4446e-04,  3.8254e-03, -4.8086e-03,  2.7346e-04,\n",
       "          3.0615e-03, -1.8756e-03,  1.6231e-04, -4.2284e-03, -5.1937e-03,\n",
       "         -2.9519e-03,  2.3717e-03,  1.1774e-03,  8.6142e-04, -5.1505e-03,\n",
       "         -6.5076e-07],\n",
       "        [ 4.9607e-03, -4.1167e-03,  2.1197e-02, -2.6602e-02,  1.5061e-03,\n",
       "          1.6967e-02, -1.0382e-02,  8.7208e-04, -2.3403e-02, -2.8751e-02,\n",
       "         -1.6338e-02,  1.3137e-02,  6.5207e-03,  4.7627e-03, -2.8504e-02,\n",
       "         -1.7984e-05],\n",
       "        [-1.1626e-02,  9.6223e-03, -4.9606e-02,  6.2195e-02, -3.5115e-03,\n",
       "         -3.9710e-02,  2.4280e-02, -2.0024e-03,  5.4731e-02,  6.7246e-02,\n",
       "          3.8207e-02, -3.0738e-02, -1.5255e-02, -1.1131e-02,  6.6656e-02,\n",
       "          6.2336e-05],\n",
       "        [ 2.3787e-02, -2.0294e-02,  1.0318e-01, -1.3079e-01,  7.6103e-03,\n",
       "          8.2512e-02, -5.0877e-02,  5.0714e-03, -1.1473e-01, -1.4079e-01,\n",
       "         -8.0113e-02,  6.4093e-02,  3.1838e-02,  2.3504e-02, -1.3983e-01,\n",
       "          3.4755e-04],\n",
       "        [-3.2886e-03,  2.7483e-03, -1.4105e-02,  1.7748e-02, -1.0119e-03,\n",
       "         -1.1288e-02,  6.9204e-03, -6.0901e-04,  1.5602e-02,  1.9161e-02,\n",
       "          1.0892e-02, -8.7472e-03, -4.3425e-03, -3.1804e-03,  1.9005e-02,\n",
       "         -3.1383e-06],\n",
       "        [ 3.8198e-03, -3.1815e-03,  1.6354e-02, -2.0552e-02,  1.1678e-03,\n",
       "          1.3089e-02, -8.0172e-03,  6.9018e-04, -1.8073e-02, -2.2200e-02,\n",
       "         -1.2617e-02,  1.0139e-02,  5.0330e-03,  3.6813e-03, -2.2015e-02,\n",
       "         -4.7469e-06],\n",
       "        [ 4.6904e-03, -3.9020e-03,  2.0069e-02, -2.5209e-02,  1.4307e-03,\n",
       "          1.6062e-02, -9.8351e-03,  8.3992e-04, -2.2171e-02, -2.7235e-02,\n",
       "         -1.5478e-02,  1.2440e-02,  6.1752e-03,  4.5146e-03, -2.7005e-02,\n",
       "         -9.5168e-06],\n",
       "        [ 1.9149e-02, -1.6033e-02,  8.2219e-02, -1.0352e-01,  5.9134e-03,\n",
       "          6.5791e-02, -4.0357e-02,  3.5947e-03, -9.0985e-02, -1.1173e-01,\n",
       "         -6.3521e-02,  5.0994e-02,  2.5317e-02,  1.8556e-02, -1.1084e-01,\n",
       "          4.1910e-05],\n",
       "        [ 4.4953e-03, -3.7465e-03,  1.9253e-02, -2.4200e-02,  1.3760e-03,\n",
       "          1.5408e-02, -9.4395e-03,  8.1590e-04, -2.1280e-02, -2.6138e-02,\n",
       "         -1.4856e-02,  1.1937e-02,  5.9254e-03,  4.3350e-03, -2.5921e-02,\n",
       "         -3.7983e-06],\n",
       "        [-1.1878e-03,  9.8407e-04, -5.0709e-03,  6.3601e-03, -3.5945e-04,\n",
       "         -4.0591e-03,  2.4826e-03, -2.0614e-04,  5.5962e-03,  6.8755e-03,\n",
       "          3.9067e-03, -3.1424e-03, -1.5596e-03, -1.1384e-03,  6.8157e-03,\n",
       "          5.6084e-06],\n",
       "        [ 1.9093e-02, -1.5949e-02,  8.1874e-02, -1.0300e-01,  5.8696e-03,\n",
       "          6.5520e-02, -4.0164e-02,  3.5235e-03, -9.0548e-02, -1.1121e-01,\n",
       "         -6.3215e-02,  5.0770e-02,  2.5205e-02,  1.8456e-02, -1.1030e-01,\n",
       "          1.2205e-05],\n",
       "        [-4.8714e-03,  4.0656e-03, -2.0879e-02,  2.6258e-02, -1.4951e-03,\n",
       "         -1.6709e-02,  1.0240e-02, -8.9333e-04,  2.3086e-02,  2.8355e-02,\n",
       "          1.6117e-02, -1.2946e-02, -6.4270e-03, -4.7046e-03,  2.8122e-02,\n",
       "         -3.5954e-07],\n",
       "        [ 6.4011e-03, -5.2991e-03,  2.7316e-02, -3.4251e-02,  1.9343e-03,\n",
       "          2.1866e-02, -1.3371e-02,  1.1045e-03, -3.0140e-02, -3.7031e-02,\n",
       "         -2.1040e-02,  1.6926e-02,  8.4007e-03,  6.1300e-03, -3.6707e-02,\n",
       "         -3.3353e-05],\n",
       "        [ 2.5201e-03, -2.0877e-03,  1.0758e-02, -1.3493e-02,  7.6257e-04,\n",
       "          8.6118e-03, -5.2670e-03,  4.3726e-04, -1.1873e-02, -1.4587e-02,\n",
       "         -8.2882e-03,  6.6668e-03,  3.3088e-03,  2.4151e-03, -1.4460e-02,\n",
       "         -1.1943e-05],\n",
       "        [ 1.3273e-03, -1.1031e-03,  5.6759e-03, -7.1271e-03,  4.0410e-04,\n",
       "          4.5430e-03, -2.7810e-03,  2.3597e-04, -6.2691e-03, -7.7012e-03,\n",
       "         -4.3765e-03,  3.5182e-03,  1.7463e-03,  1.2762e-03, -7.6357e-03,\n",
       "         -3.5247e-06],\n",
       "        [-3.8757e-03,  3.2165e-03, -1.6561e-02,  2.0785e-02, -1.1768e-03,\n",
       "         -1.3256e-02,  8.1115e-03, -6.8160e-04,  1.8285e-02,  2.2464e-02,\n",
       "          1.2765e-02, -1.0264e-02, -5.0947e-03, -3.7212e-03,  2.2271e-02,\n",
       "          1.3925e-05],\n",
       "        [ 6.8812e-03, -5.7186e-03,  2.9426e-02, -3.6948e-02,  2.0948e-03,\n",
       "          2.3552e-02, -1.4417e-02,  1.2227e-03, -3.2500e-02, -3.9925e-02,\n",
       "         -2.2689e-02,  1.8239e-02,  9.0534e-03,  6.6162e-03, -3.9585e-02,\n",
       "         -1.8580e-05],\n",
       "        [-2.0882e-02,  1.7397e-02, -8.9417e-02,  1.1238e-01, -6.3874e-03,\n",
       "         -7.1563e-02,  4.3837e-02, -3.7802e-03,  9.8823e-02,  1.2139e-01,\n",
       "          6.8991e-02, -5.5436e-02, -2.7519e-02, -2.0130e-02,  1.2037e-01,\n",
       "          2.2477e-05]])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jac[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "d7557240",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.0733,  0.2283,  0.0228, -0.2164, -0.0393, -0.0296, -0.1130,  0.0097,\n",
       "          0.0356, -0.1698, -0.0225, -0.0267, -0.2118,  0.0152, -0.0715,  0.0792]])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jac[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "5175d7a1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "20"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "numpy.prod(outputs[0].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "07b624b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here's a sort of trick: scale out the batch size by n_ouputs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "6836f9a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "params = d.parameters()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "fb366597",
   "metadata": {},
   "outputs": [],
   "source": [
    "shapes = (p.size() for p in params)\n",
    "n_params = numpy.sum([numpy.prod(s) for s in shapes])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "e48b48d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here's the stupid way:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 300,
   "id": "b331b7a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def jacobian_reverse(inputs, n_params, model):\n",
    "    '''\n",
    "    This computes a backward pass once per input in a way that builds up the full\n",
    "    jacobian matrix (ninputs x nparameters) one input at a time.\n",
    "    \n",
    "    For deep networks / many parameters, and for just a few inputs, this is more efficient.\n",
    "    '''\n",
    "    n_walkers = inputs.size()[0]\n",
    "    jac_output = torch.zeros((n_walkers, n_params)) \n",
    "    o = model(inputs)\n",
    "\n",
    "    for i_walker in range(n_walkers):\n",
    "        grad_outputs = torch.zeros_like(o)\n",
    "        grad_outputs[i_walker] = 1.0\n",
    "        single_jac = torch.autograd.grad(o, model.parameters(), retain_graph=True, grad_outputs=grad_outputs)\n",
    "        flattened_line = torch.cat([s.flatten() for s in single_jac])\n",
    "        jac_output[i_walker, :] = flattened_line\n",
    "    return jac_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 323,
   "id": "8346ba18",
   "metadata": {},
   "outputs": [],
   "source": [
    "j_bkwd = jacobian_reverse(x, n_params, d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 341,
   "id": "91518c87",
   "metadata": {},
   "outputs": [],
   "source": [
    "def jacobian_forward(inputs, n_params, model):\n",
    "    '''\n",
    "    https://j-towns.github.io/2017/06/12/A-new-trick.html\n",
    "    \n",
    "    Based on the above trick, this uses two backward passes to arrive at a forward-mode gradient.\n",
    "    The jacobian then is constructed one parameter at a time (instead of one input at a time).\n",
    "    For smaller networks, or when n_params << n_inputs, this should be more efficient.\n",
    "    \n",
    "    '''\n",
    "    network_output = model(inputs)\n",
    "    n_walkers = inputs.size()[0]\n",
    "    jac_output = torch.zeros((n_walkers, n_params)) \n",
    "    \n",
    "    # Loop over layers in the model.\n",
    "    # Keep track of the column of the jacobian:\n",
    "    running_column_index = 0\n",
    "    for i_layer, layer in enumerate(model.parameters()):\n",
    "        \n",
    "        # How many parameters in this layer?\n",
    "        n_params_local = numpy.prod(layer.shape)\n",
    "        \n",
    "        for i_weight in range(n_params_local):\n",
    "            v = torch.ones_like(network_output, requires_grad=True)\n",
    "            u = torch.zeros_like(layer).flatten()\n",
    "            # Need to set the i^th index to 1.0, but it's trick to do so without reshaping:\n",
    "            u[i_weight] = 1.0\n",
    "            u = u.reshape(layer.shape)\n",
    "            # First backward pass:\n",
    "            vjp = torch.autograd.grad(network_output, layer, grad_outputs=v, create_graph=True)[0]\n",
    "            # Second backward pass:\n",
    "            output = torch.autograd.grad(vjp, v, grad_outputs=u)[0]\n",
    "            jac_output[:,running_column_index] = output\n",
    "            running_column_index += 1\n",
    "    return jac_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 342,
   "id": "fd455210",
   "metadata": {},
   "outputs": [],
   "source": [
    "j_fwd = jacobian_forward(x, n_params, d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 343,
   "id": "38318603",
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "The size of tensor a (1000) must match the size of tensor b (20) at non-singleton dimension 0",
     "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_15141/3815104110.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmax\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mabs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mj_fwd\u001b[0m \u001b[0;34m-\u001b[0m \u001b[0mj_bkwd\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[0m",
      "\u001b[0;31mRuntimeError\u001b[0m: The size of tensor a (1000) must match the size of tensor b (20) at non-singleton dimension 0"
     ]
    }
   ],
   "source": [
    "torch.max(torch.abs(j_fwd - j_bkwd))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 344,
   "id": "f76675bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Which is more efficient?\n",
    "n_walkers = 20\n",
    "n_particles = 2\n",
    "n_dim = 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 345,
   "id": "e537262a",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.rand((n_walkers, n_particles, n_dim), requires_grad=True, dtype=torch.float64)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 346,
   "id": "f27cb233",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([20, 2, 3])\n",
      "1088\n"
     ]
    }
   ],
   "source": [
    "print(x.shape)\n",
    "print(n_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 347,
   "id": "62da0acf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "272 ms ± 11.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%timeit j_fwd = jacobian_forward(x, n_params, d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 348,
   "id": "08ee2485",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4.25 ms ± 211 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit j_bkwd = jacobian_reverse(x, n_params, d)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c9bd69f",
   "metadata": {},
   "source": [
    "For small numbers of walkers, reverse is better.  What about comparable sizes?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 349,
   "id": "01f560c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_walkers = 1000\n",
    "n_particles = 2\n",
    "n_dim = 3\n",
    "x = torch.rand((n_walkers, n_particles, n_dim), requires_grad=True, dtype=torch.float64)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 350,
   "id": "08144078",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "835 ms ± 27.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%timeit j_fwd = jacobian_forward(x, n_params, d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 351,
   "id": "2e329923",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "437 ms ± 3.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%timeit j_bkwd = jacobian_reverse(x, n_params, d)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f4409a1",
   "metadata": {},
   "source": [
    "Here, it took about twice as long for the forward mode as backward mode - makes sense, its 2x passes per parameter and there are similar sized parameters as inputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 352,
   "id": "946d64db",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_walkers = 10000\n",
    "n_particles = 2\n",
    "n_dim = 3\n",
    "x = torch.rand((n_walkers, n_particles, n_dim), requires_grad=True, dtype=torch.float64)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 355,
   "id": "3f5b8ddd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "9.73 s ± 198 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%timeit j_fwd = jacobian_forward(x, n_params, d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 356,
   "id": "c901930e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "45.8 s ± 1.12 s per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%timeit j_bkwd = jacobian_reverse(x, n_params, d)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "966f799d",
   "metadata": {},
   "source": [
    "At very high number of walkers, which is more efficient for memory usage, the forward mode jacobian is better!  In particular, there is only a little overhead introduced "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c1c14c7",
   "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
}