notebooks/Hijacking-Autograd-for-LRP.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hijacking autograd for LRP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "import numpy as np\n",
    "\n",
    "import torch\n",
    "import torch_scatter\n",
    "import torchgraphs as tg\n",
    "\n",
    "np.set_printoptions(formatter={'float_kind': '{:5.2f}'.format, 'int_kind': '{:5d}'.format}, linewidth=150)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Add"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 1.00  2.00  3.00  1.00] + [ 0.00  1.00  2.00  4.00] ----> [ 1.00  3.00  5.00  5.00]\n",
      "[ 1.00  1.00  1.00  1.00]   [ 1.00  1.00  1.00  1.00] <---- [ 1.00  1.00  1.00  1.00]\n"
     ]
    }
   ],
   "source": [
    "a = torch.tensor([1, 2, 3, 1], dtype=torch.float, requires_grad=True)\n",
    "b = torch.tensor([0, 1, 2, 4], dtype=torch.float, requires_grad=True)\n",
    "c = a + b\n",
    "\n",
    "grad_out = torch.ones_like(c) * (c != 0).float()\n",
    "c.backward(grad_out)\n",
    "\n",
    "print(a.detach().numpy(), '+', b.detach().numpy(), '---->', c.detach().numpy())\n",
    "print(a.grad.numpy(), ' ', b.grad.numpy(), '<----', grad_out.detach().numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AddRelevance(torch.autograd.Function):\n",
    "    @staticmethod\n",
    "    def forward(ctx, a, b):\n",
    "        out = a + b\n",
    "        ctx.save_for_backward(a, b, out)\n",
    "        return out\n",
    "\n",
    "    @staticmethod\n",
    "    def backward(ctx, rel_out):\n",
    "        a, b, out = ctx.saved_tensors\n",
    "        if ((out == 0) & (rel_out > 0)).any():\n",
    "            warnings.warn('Relevance that is propagated back through an output of 0 will be lost')\n",
    "        rel_a = torch.where(out != 0, rel_out * a / out, out.new_tensor(0))\n",
    "        rel_b = torch.where(out != 0, rel_out * b / out, out.new_tensor(0))\n",
    "        return rel_a, rel_b\n",
    "    \n",
    "def add(a, b):\n",
    "    return AddRelevance.apply(a, b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 1.00  2.00  3.00  1.00] + [ 0.00  1.00  2.00  4.00] ----> [ 1.00  3.00  5.00  5.00]\n",
      "[ 1.00  0.67  0.60  0.20]   [ 0.00  0.33  0.40  0.80] <---- [ 1.00  1.00  1.00  1.00]\n"
     ]
    }
   ],
   "source": [
    "a = torch.tensor([1, 2, 3, 1], dtype=torch.float, requires_grad=True)\n",
    "b = torch.tensor([0, 1, 2, 4], dtype=torch.float, requires_grad=True)\n",
    "c = add(a, b)\n",
    "\n",
    "rel_out = torch.ones_like(c) * (c != 0).float()\n",
    "c.backward(torch.ones_like(rel_out))\n",
    "\n",
    "print(a.detach().numpy(), '+', b.detach().numpy(), '---->', c.detach().numpy())\n",
    "print(a.grad.numpy(), ' ', b.grad.numpy(), '<----', rel_out.numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Add"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 1.00  2.00  3.00  1.00]\n",
      " [-6.00 -2.00 -1.00 -1.00]] ----> [-5.00  0.00  2.00  0.00]\n",
      "[[ 1.00  0.00  1.00  0.00]\n",
      " [ 1.00  0.00  1.00  0.00]] <---- [ 1.00  0.00  1.00  0.00]\n"
     ]
    }
   ],
   "source": [
    "a = torch.tensor([[1, 2, 3, 1], [-6, -2, -1, -1]], dtype=torch.float, requires_grad=True)\n",
    "b = torch.sum(a, dim=0)\n",
    "\n",
    "grad_out = torch.ones_like(b) * (b != 0).float()\n",
    "b.backward(grad_out)\n",
    "\n",
    "print(a.detach().numpy(), '---->', b.detach().numpy())\n",
    "print(a.grad.numpy(), '<----', grad_out.detach().numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SumPooling(torch.autograd.Function):\n",
    "    @staticmethod\n",
    "    def forward(ctx, src, dim, keepdim):\n",
    "        out = torch.sum(src, dim=dim, keepdim=keepdim)\n",
    "        ctx.dim = dim\n",
    "        ctx.keepdim = keepdim\n",
    "        ctx.save_for_backward(src, out)\n",
    "        return out\n",
    "\n",
    "    @staticmethod\n",
    "    def backward(ctx, rel_out):\n",
    "        src, out = ctx.saved_tensors\n",
    "        if ((out == 0) & (rel_out > 0)).any():\n",
    "            warnings.warn('Relevance that is propagated back through an output of 0 will be lost')\n",
    "        rel_out = torch.where(out != 0, rel_out / out, out.new_tensor(0))\n",
    "        if not ctx.keepdim and ctx.dim is not None:\n",
    "            rel_out.unsqueeze_(ctx.dim)\n",
    "        return rel_out * src, None, None\n",
    "\n",
    "\n",
    "def sum(tensor, dim=None, keepdim=False):\n",
    "    return SumPooling.apply(tensor, dim, keepdim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 1.00  2.00  3.00  1.00]\n",
      " [-6.00 -2.00 -1.00 -1.00]] ----> [-5.00  0.00  2.00  0.00]\n",
      "[[-0.20  0.00  1.50  0.00]\n",
      " [ 1.20 -0.00 -0.50 -0.00]] <---- [ 1.00  0.00  1.00  0.00]\n"
     ]
    }
   ],
   "source": [
    "a = torch.tensor([[1, 2, 3, 1], [-6, -2, -1, -1]], dtype=torch.float, requires_grad=True)\n",
    "b = sum(a, dim=0)\n",
    "\n",
    "rel_out = torch.ones_like(b) * (b != 0).float()\n",
    "b.backward(rel_out)\n",
    "\n",
    "print(a.detach().numpy(), '---->', b.detach().numpy())\n",
    "print(a.grad.numpy(), '<----', rel_out.detach().numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 1.00  2.00  3.00  1.00]\n",
      " [-6.00 -2.00 -1.00 -1.00]] ----> [[-5.00  0.00  2.00  0.00]]\n",
      "[[-0.20  0.00  1.50  0.00]\n",
      " [ 1.20 -0.00 -0.50 -0.00]] <---- [[ 1.00  0.00  1.00  0.00]]\n"
     ]
    }
   ],
   "source": [
    "a = torch.tensor([[1, 2, 3, 1], [-6, -2, -1, -1]], dtype=torch.float, requires_grad=True)\n",
    "b = sum(a, dim=0, keepdim=True)\n",
    "\n",
    "rel_out = torch.ones_like(b) * (b != 0).float()\n",
    "b.backward(rel_out)\n",
    "\n",
    "print(a.detach().numpy(), '---->', b.detach().numpy())\n",
    "print(a.grad.numpy(), '<----', rel_out.detach().numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 1.00  2.00  3.00  1.00]\n",
      " [-6.00 -2.00 -1.00 -1.00]] ----> [ 7.00 -10.00]\n",
      "[[ 0.14  0.29  0.43  0.14]\n",
      " [ 0.60  0.20  0.10  0.10]] <---- [ 1.00  1.00]\n"
     ]
    }
   ],
   "source": [
    "a = torch.tensor([[1, 2, 3, 1], [-6, -2, -1, -1]], dtype=torch.float, requires_grad=True)\n",
    "b = sum(a, dim=1)\n",
    "\n",
    "rel_out = torch.ones_like(b) * (b != 0).float()\n",
    "b.backward(rel_out)\n",
    "\n",
    "print(a.detach().numpy(), '---->', b.detach().numpy())\n",
    "print(a.grad.numpy(), '<----', rel_out.detach().numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Scatter Add"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[    0     0     0     0     1     1     2     2     2]\n",
      "[ 1.00  1.00  1.00  1.00  2.00  6.00  1.00  1.00 -2.00] ----> [ 4.00  8.00  0.00]\n",
      "[ 1.00  1.00  1.00  1.00  1.00  1.00  0.00  0.00  0.00] <---- [ 1.00  1.00  0.00]\n"
     ]
    }
   ],
   "source": [
    "a = torch.tensor(\n",
    "    [1, 1, 1, 1, 2, 6, 1, 1, -2], dtype=torch.float, requires_grad=True)\n",
    "a_idx = torch.tensor(\n",
    "    [0, 0, 0, 0, 1, 1, 2, 2, 2])\n",
    "a_new = torch_scatter.scatter_add(a, a_idx, dim=0)\n",
    "\n",
    "grad_out = torch.ones_like(a_new) * (a_new != 0).float()\n",
    "a_new.backward(grad_out)\n",
    "\n",
    "print(a_idx.numpy())\n",
    "print(a.detach().numpy(), '---->', a_new.detach().numpy())\n",
    "print(a.grad.numpy(), '<----', grad_out.detach().numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ScatterAddRelevance(torch.autograd.Function):\n",
    "    @staticmethod\n",
    "    def forward(ctx, src, idx, dim, dim_size, fill_value):\n",
    "        out = torch_scatter.scatter_add(src, idx, dim, None, dim_size, fill_value)\n",
    "        ctx.dim = dim\n",
    "        ctx.save_for_backward(src, idx, out)\n",
    "        return out\n",
    "\n",
    "    @staticmethod\n",
    "    def backward(ctx, rel_out):\n",
    "        src, idx, out = ctx.saved_tensors\n",
    "        if ((out == 0) & (rel_out > 0)).any():\n",
    "            warnings.warn('Relevance that is propagated back through an output of 0 will be lost')\n",
    "        rel_out = torch.where(out != 0, rel_out / out, out.new_tensor(0))\n",
    "        rel_src = torch.index_select(rel_out, ctx.dim, idx) * src\n",
    "        return rel_src, None, None, None, None\n",
    "    \n",
    "def scatter_add(src, index, dim=-1, out=None, dim_size=None, fill_value=0):\n",
    "    return ScatterAddRelevance.apply(src, index, dim, dim_size, fill_value)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 1.00  1.00  1.00  1.00  2.00  6.00  1.00  1.00 -2.00] ----> [ 4.00  8.00  0.00]\n",
      "[ 0.25  0.25  0.25  0.25  0.25  0.75  0.00  0.00 -0.00] <---- [ 1.00  1.00  0.00]\n"
     ]
    }
   ],
   "source": [
    "b = torch.tensor(\n",
    "    [1, 1, 1, 1, 2, 6, 1, 1, -2], dtype=torch.float, requires_grad=True)\n",
    "b_idx = torch.tensor(\n",
    "    [0, 0, 0, 0, 1, 1, 2, 2, 2])\n",
    "b_new = scatter_add(b, b_idx, dim=0)\n",
    "\n",
    "rel_out = torch.ones_like(b_new) * (b_new != 0).float()\n",
    "b_new.backward(rel_out)\n",
    "\n",
    "print(b.detach().numpy(), '---->', b_new.detach().numpy())\n",
    "print(b.grad.numpy(), '<----', rel_out.numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Scatter Mean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[    0     0     0     0     1     1     2     2     2]\n",
      "[ 1.00  1.00  1.00  1.00  2.00  8.00  1.00  1.00 -2.00] ----> [ 1.00  5.00  0.00]\n",
      "[ 0.25  0.25  0.25  0.25  0.50  0.50  0.00  0.00  0.00] <---- [ 1.00  1.00  0.00]\n"
     ]
    }
   ],
   "source": [
    "a = torch.tensor(\n",
    "    [1, 1, 1, 1, 2, 8, 1, 1, -2], dtype=torch.float, requires_grad=True)\n",
    "a_idx = torch.tensor(\n",
    "    [0, 0, 0, 0, 1, 1, 2, 2, 2])\n",
    "a_new = torch_scatter.scatter_mean(a, a_idx, dim=0)\n",
    "\n",
    "grad_out = torch.full_like(a_new, 1) * (a_new != 0).float()\n",
    "a_new.backward(grad_out)\n",
    "\n",
    "print(a_idx.numpy())\n",
    "print(a.detach().numpy(), '---->', a_new.detach().numpy())\n",
    "print(a.grad.numpy(), '<----', grad_out.detach().numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ScatterMeanRelevance(torch.autograd.Function):\n",
    "    @staticmethod\n",
    "    def forward(ctx, src, idx, dim, dim_size, fill_value):\n",
    "        sums = torch_scatter.scatter_add(src, idx, dim, None, dim_size, fill_value)\n",
    "        count = torch_scatter.scatter_add(torch.ones_like(src), idx, dim, None, dim_size, fill_value=0)\n",
    "        out =  sums / count.clamp(min=1)\n",
    "        ctx.dim = dim\n",
    "        ctx.save_for_backward(src, idx, sums)\n",
    "        return out\n",
    "\n",
    "    @staticmethod\n",
    "    def backward(ctx, rel_out):\n",
    "        src, idx, sums = ctx.saved_tensors\n",
    "        if ((sums == 0) & (rel_out > 0)).any():\n",
    "            warnings.warn('Relevance that is propagated back through an output of 0 will be lost')\n",
    "        rel_out = torch.where(sums != 0, rel_out / sums, sums.new_tensor(0))\n",
    "        rel_src = torch.index_select(rel_out, ctx.dim, idx) * src\n",
    "        return rel_src, None, None, None, None\n",
    "\n",
    "def scatter_mean(src, index, dim=-1, dim_size=None, fill_value=0):\n",
    "    return ScatterMeanRelevance.apply(src, index, dim, dim_size, fill_value)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 1.00  1.00  1.00  1.00  2.00  8.00  1.00  1.00 -2.00] ----> [ 1.00  5.00  0.00]\n",
      "[ 0.25  0.25  0.25  0.25  0.20  0.80  0.00  0.00 -0.00] <---- [ 1.00  1.00  0.00]\n"
     ]
    }
   ],
   "source": [
    "b = torch.tensor(\n",
    "    [1, 1, 1, 1, 2, 8, 1, 1, -2], dtype=torch.float, requires_grad=True)\n",
    "b_idx = torch.tensor(\n",
    "    [0, 0, 0, 0, 1, 1, 2, 2, 2])\n",
    "b_new = scatter_mean(b, b_idx, dim=0)\n",
    "\n",
    "rel_out = torch.ones_like(b_new) * (b_new != 0).float()\n",
    "b_new.backward(rel_out)\n",
    "\n",
    "print(b.detach().numpy(), '---->', b_new.detach().numpy())\n",
    "print(b.grad.numpy(), '<----', rel_out.numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Scatter Max"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[    0     0     0     0     1     1     2     2     2]\n",
      "[ 1.00  1.00  1.00  1.00  2.00  8.00  1.00  1.00 -2.00] ----> [ 1.00  8.00  1.00]\n",
      "[ 0.00  0.00  0.00  1.00  0.00  1.00  0.00  1.00  0.00] <---- [ 1.00  1.00  1.00]\n"
     ]
    }
   ],
   "source": [
    "a = torch.tensor(\n",
    "    [1, 1, 1, 1, 2, 8, 1, 1, -2], dtype=torch.float, requires_grad=True)\n",
    "a_idx = torch.tensor(\n",
    "    [0, 0, 0, 0, 1, 1, 2, 2, 2])\n",
    "a_new = torch_scatter.scatter_max(a, a_idx, dim=0)[0]\n",
    "\n",
    "grad_out = torch.full_like(a_new, 1) * (a_new != 0).float()\n",
    "a_new.backward(grad_out)\n",
    "\n",
    "print(a_idx.numpy())\n",
    "print(a.detach().numpy(), '---->', a_new.detach().numpy())\n",
    "print(a.grad.numpy(), '<----', grad_out.detach().numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ScatterMaxRelevance(torch.autograd.Function):\n",
    "    @staticmethod\n",
    "    def forward(ctx, src, idx, dim, dim_size, fill_value):\n",
    "        out, idx_maxes = torch_scatter.scatter_max(src, idx, dim=dim, dim_size=dim_size, fill_value=fill_value)\n",
    "        ctx.dim = dim\n",
    "        ctx.dim_size = src.shape[dim]\n",
    "        ctx.save_for_backward(idx, out, idx_maxes)\n",
    "        return out, idx_maxes\n",
    "\n",
    "    @staticmethod\n",
    "    def backward(ctx, rel_out, rel_idx_maxes):\n",
    "        idx, out, idx_maxes = ctx.saved_tensors\n",
    "        if ((out == 0) & (rel_out > 0)).any():\n",
    "            warnings.warn('Relevance that is propagated back through an output of 0 will be lost')\n",
    "        rel_out = torch.where(out != 0, rel_out, out.new_tensor(0))\n",
    "        \n",
    "        # Where idx_maxes==-1 set idx=0 so that the indexes are valid for scatter_add\n",
    "        # The corresponding relevance should already be 0, but set it relevance=0 to be sure\n",
    "        rel_out = torch.where(idx_maxes != -1, rel_out, torch.zeros_like(rel_out))\n",
    "        idx_maxes = torch.where(idx_maxes != -1, idx_maxes, torch.zeros_like(idx_maxes))\n",
    "\n",
    "        rel_src = torch_scatter.scatter_add(rel_out, idx_maxes, dim=ctx.dim, dim_size=ctx.dim_size)\n",
    "        return rel_src, None, None, None, None\n",
    "\n",
    "def scatter_max(src, index, dim=-1, dim_size=None, fill_value=0):\n",
    "    return ScatterMaxRelevance.apply(src, index, dim, dim_size, fill_value)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 1.00  1.00  1.00  1.00  2.00  8.00  0.00  0.00  0.00] ----> [ 1.00  8.00  0.00  0.00]\n",
      "[ 0.00  0.00  0.00  3.00  0.00  3.00  0.00  0.00  0.00] <---- [ 3.00  3.00  0.00  0.00]\n"
     ]
    }
   ],
   "source": [
    "b = torch.tensor(\n",
    "    [1, 1, 1, 1, 2, 8, 0, 0, 0], dtype=torch.float, requires_grad=True)\n",
    "b_idx = torch.tensor(\n",
    "    [0, 0, 0, 0, 1, 1, 3, 3, 3])\n",
    "b_new = scatter_max(b, b_idx, dim=0)[0]\n",
    "\n",
    "rel_out = torch.full_like(b_new, 3) * (b_new != 0).float()\n",
    "b_new.backward(rel_out)\n",
    "\n",
    "print(b.detach().numpy(), '---->', b_new.detach().numpy())\n",
    "print(b.grad.numpy(), '<----', rel_out.numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Linear layer with epsilon rule"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 4.00  8.00  0.00]] ----> [[20.00  1.00]]\n",
      "[[ 1.00  2.00  0.00]] <---- [[ 1.00  1.00]]\n"
     ]
    }
   ],
   "source": [
    "weight = torch.tensor([\n",
    "    [1, 2, -1],\n",
    "    [0, 0, +1],\n",
    "], dtype=torch.float)\n",
    "bias = torch.tensor([0, 1], dtype=torch.float)\n",
    "x = torch.tensor([[4, 8, 0]], dtype=torch.float, requires_grad=True)\n",
    "y = x @ weight.t() + bias\n",
    "\n",
    "grad_out = torch.ones_like(y) * (y != 0).float()\n",
    "y.backward(grad_out)\n",
    "\n",
    "print(x.detach().numpy(), '---->', y.detach().numpy())\n",
    "print(x.grad.numpy(), '<----', grad_out.detach().numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LinearEpsilonRelevance(torch.autograd.Function):\n",
    "    eps = 1e-16\n",
    "    \n",
    "    @staticmethod\n",
    "    def forward(ctx, input, weight, bias):\n",
    "        Z = weight.t()[None, :, :] * input[:, :, None]\n",
    "        Zs = Z.sum(dim=1, keepdim=True)\n",
    "        if bias is not None:\n",
    "            Zs += bias[None, None, :]\n",
    "        ctx.save_for_backward(Z, Zs)\n",
    "        return Zs.squeeze(dim=1)\n",
    "\n",
    "    @staticmethod\n",
    "    def backward(ctx, rel_out):\n",
    "        Z, Zs = ctx.saved_tensors\n",
    "        eps = rel_out.new_tensor(LinearEpsilonRelevance.eps)\n",
    "        Zs += torch.where(Zs >= 0, eps, -eps)\n",
    "        return (rel_out[:, None, :] * Z / Zs).sum(dim=2), None, None\n",
    "    \n",
    "def linear_eps(input, weight, bias=None):\n",
    "    return LinearEpsilonRelevance.apply(input, weight, bias)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "No bias -> conservation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "W\n",
      "[[ 1.00  2.00 -1.00]\n",
      " [ 0.00  0.00  1.00]]\n",
      "\n",
      "[[ 4.00  8.00  0.00]] ----> [[20.00  0.00]]\n",
      "[[ 0.20  0.80  0.00]] <---- [[ 1.00  0.00]]\n"
     ]
    }
   ],
   "source": [
    "weight = torch.tensor([\n",
    "    [1, 2, -1],\n",
    "    [0, 0, +1],\n",
    "], dtype=torch.float)\n",
    "bias = None\n",
    "print('W', weight.numpy(), sep='\\n', end='\\n\\n')\n",
    "\n",
    "x = torch.tensor([\n",
    "    [4, 8, 0]\n",
    "], dtype=torch.float, requires_grad=True)\n",
    "\n",
    "y = linear_eps(x, weight, bias)\n",
    "rel_out = torch.ones_like(y) * (y != 0).float()\n",
    "y.backward(rel_out)\n",
    "\n",
    "print(x.detach().numpy(), '---->', y.detach().numpy())\n",
    "print(x.grad.numpy().round(2), '<----', rel_out.numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Bias absorbs relevance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "W\n",
      "[[ 1.00  2.00 -1.00]\n",
      " [ 0.00  0.00  1.00]]\n",
      "b\n",
      "[ 0.00  1.00]\n",
      "\n",
      "[[ 4.00  8.00  0.00]] ----> [[20.00  1.00]]\n",
      "[[ 0.20  0.80  0.00]] <---- [[ 1.00  1.00]]\n"
     ]
    }
   ],
   "source": [
    "weight = torch.tensor([\n",
    "    [1, 2, -1],\n",
    "    [0, 0, +1],\n",
    "], dtype=torch.float)\n",
    "bias = torch.tensor([0, 1], dtype=torch.float)\n",
    "print('W', weight.numpy(), 'b', bias.numpy(), sep='\\n', end='\\n\\n')\n",
    "\n",
    "x = torch.tensor([\n",
    "    [4, 8, 0]\n",
    "], dtype=torch.float, requires_grad=True)\n",
    "\n",
    "y = linear_eps(x, weight, bias)\n",
    "rel_out = torch.ones_like(y) * (y != 0).float()\n",
    "y.backward(rel_out)\n",
    "\n",
    "print(x.detach().numpy(), '---->', y.detach().numpy())\n",
    "print(x.grad.numpy().round(2), '<----', rel_out.numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But in a weird way?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "W\n",
      "[[ 1.00  2.00 -1.00]\n",
      " [ 0.00  0.00  1.00]]\n",
      "b\n",
      "[ 0.00 100.00]\n",
      "\n",
      "[[ 4.00  8.00  0.00]] ----> [[20.00 100.00]]\n",
      "[[ 0.20  0.80  0.00]] <---- [[ 1.00  1.00]]\n"
     ]
    }
   ],
   "source": [
    "weight = torch.tensor([\n",
    "    [1, 2, -1],\n",
    "    [0, 0, +1],\n",
    "], dtype=torch.float)\n",
    "bias = torch.tensor([0, 100], dtype=torch.float)\n",
    "print('W', weight.numpy(), 'b', bias.numpy(), sep='\\n', end='\\n\\n')\n",
    "\n",
    "x = torch.tensor([\n",
    "    [4, 8, 0]\n",
    "], dtype=torch.float, requires_grad=True)\n",
    "\n",
    "y = linear_eps(x, weight, bias)\n",
    "rel_out = torch.ones_like(y) * (y != 0).float()\n",
    "y.backward(rel_out)\n",
    "\n",
    "print(x.detach().numpy(), '---->', y.detach().numpy())\n",
    "print(x.grad.numpy().round(2), '<----', rel_out.numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Zeros in the input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "W\n",
      "[[ 1.00 -1.00]]\n",
      "\n",
      "[[ 4.00  0.00]] ----> [[ 4.00]]\n",
      "[[ 1.00  0.00]] <---- [[ 1.00]]\n"
     ]
    }
   ],
   "source": [
    "weight = torch.tensor([\n",
    "    [+1, -1],\n",
    "], dtype=torch.float)\n",
    "bias = None\n",
    "print('W', weight.numpy(), sep='\\n', end='\\n\\n')\n",
    "\n",
    "x = torch.tensor([\n",
    "    [4, 0]\n",
    "], dtype=torch.float, requires_grad=True)\n",
    "\n",
    "y = linear_eps(x, weight, bias)\n",
    "rel_out = torch.ones_like(y) * (y != 0).float()\n",
    "y.backward(rel_out)\n",
    "\n",
    "print(x.detach().numpy(), '---->', y.detach().numpy())\n",
    "print(x.grad.numpy().round(2), '<----', rel_out.numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Zeros in the weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "W\n",
      "[[ 1.00  0.00  2.00]]\n",
      "\n",
      "[[ 4.00  1.00  3.00]] ----> [[10.00]]\n",
      "[[ 0.40  0.00  0.60]] <---- [[ 1.00]]\n"
     ]
    }
   ],
   "source": [
    "weight = torch.tensor([\n",
    "    [1, 0, 2],\n",
    "], dtype=torch.float)\n",
    "bias = None\n",
    "print('W', weight.numpy(), sep='\\n', end='\\n\\n')\n",
    "\n",
    "x = torch.tensor([\n",
    "    [4, 1, 3]\n",
    "], dtype=torch.float, requires_grad=True)\n",
    "\n",
    "y = linear_eps(x, weight, bias)\n",
    "rel_out = torch.ones_like(y) * (y != 0).float()\n",
    "y.backward(rel_out)\n",
    "\n",
    "print(x.detach().numpy(), '---->', y.detach().numpy())\n",
    "print(x.grad.numpy().round(2), '<----', rel_out.numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Putting things together"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 1.00  2.00  3.00  1.00  0.00]] + [[ 0.00  1.00  2.00  4.00  5.00]] ----> [[ 1.00  3.00  5.00  5.00  5.00]]\n",
      "\n",
      "[    0     0     1     1     2]\n",
      "[[ 1.00  3.00  5.00  5.00  5.00]] ----> [[ 4.00 10.00  5.00]]\n",
      "\n",
      "[[ 1.00  0.00  2.00]\n",
      " [ 0.00  5.00  1.00]]\n",
      "[[ 4.00 10.00  5.00]] ----> [[14.00 55.00]]\n",
      "\n",
      "[[ 0.00  1.00 -1.00]\n",
      " [-4.00 -2.00 -3.00]]\n",
      "[[ 4.00 10.00  5.00]] ----> [[ 5.00 -51.00]]\n",
      "\n",
      "[[14.00 55.00]] + [[ 5.00 -51.00]] ----> [[19.00  4.00]]\n",
      "\n",
      "[[-0.95 -1.89  2.41  0.80 -0.00]]   [[-0.00 -0.95  1.61  3.21 -2.24]] <---- [[ 1.00  1.00]]\n",
      "   0.36842078 \t\t\t  +    1.6315787 \t\t\t==     2.0\n"
     ]
    }
   ],
   "source": [
    "a = torch.tensor([[1, 2, 3, 1, 0]], dtype=torch.float, requires_grad=True)\n",
    "b = torch.tensor([[0, 1, 2, 4, 5]], dtype=torch.float, requires_grad=True)\n",
    "c = add(a, b)\n",
    "\n",
    "print(a.detach().numpy(), '+', b.detach().numpy(), '---->', c.detach().numpy(), end='\\n\\n')\n",
    "\n",
    "idx = torch.tensor([0, 0, 1, 1, 2])\n",
    "d = scatter_add(c, idx, dim=1)\n",
    "\n",
    "print(idx.numpy())\n",
    "print(c.detach().numpy(), '---->', d.detach().numpy(), end='\\n\\n')\n",
    "\n",
    "weight_u = torch.tensor([\n",
    "    [1, 0, 2],\n",
    "    [0, 5, 1],\n",
    "], dtype=torch.float)\n",
    "u = linear_eps(d, weight_u, None)\n",
    "\n",
    "print(weight_u.numpy())\n",
    "print(d.detach().numpy(), '---->', u.detach().numpy(), end='\\n\\n')\n",
    "\n",
    "weight_v = torch.tensor([\n",
    "    [ 0,  1, -1],\n",
    "    [-4, -2, -3],\n",
    "], dtype=torch.float)\n",
    "v = linear_eps(d, weight_v, None)\n",
    "\n",
    "print(weight_v.numpy())\n",
    "print(d.detach().numpy(), '---->', v.detach().numpy(), end='\\n\\n')\n",
    "\n",
    "z = add(u, v)\n",
    "print(u.detach().numpy(), '+', v.detach().numpy(), '---->', z.detach().numpy(), end='\\n\\n')\n",
    "\n",
    "rel_out = torch.ones_like(z) * (z != 0).float()\n",
    "z.backward(rel_out)\n",
    "print(a.grad.numpy(), ' ', b.grad.numpy(), '<----', rel_out.numpy())\n",
    "print('  ', a.grad.numpy().sum(), '\\t\\t\\t  +   ', b.grad.numpy().sum(), '\\t\\t\\t==    ', rel_out.numpy().sum())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "- <torch.autograd.function.AddRelevanceBackward object at 0x55ee07c718c8> \n",
      "  - <torch.autograd.function.LinearEpsilonRelevanceBackward object at 0x55ee08db8b08> \n",
      "    - <torch.autograd.function.ScatterAddRelevanceBackward object at 0x55ee08acc198> \n",
      "      - <torch.autograd.function.AddRelevanceBackward object at 0x55ee08e00c48> \n",
      "        - <AccumulateGrad object at 0x7f8bfd8bac18> ([[ 1.00  2.00  3.00  1.00  0.00]] at 0x7f8bfc8909d8)\n",
      "        - <AccumulateGrad object at 0x7f8bfd8babe0> ([[ 0.00  1.00  2.00  4.00  5.00]] at 0x7f8bfd8d9678)\n",
      "      - None \n",
      "    - None \n",
      "  - <torch.autograd.function.LinearEpsilonRelevanceBackward object at 0x55ee0783b618> \n",
      "    - <torch.autograd.function.ScatterAddRelevanceBackward object at 0x55ee08acc198> \n",
      "      - <torch.autograd.function.AddRelevanceBackward object at 0x55ee08e00c48> \n",
      "        - <AccumulateGrad object at 0x7f8bfd8babe0> ([[ 1.00  2.00  3.00  1.00  0.00]] at 0x7f8bfc8909d8)\n",
      "        - <AccumulateGrad object at 0x7f8bfd8bafd0> ([[ 0.00  1.00  2.00  4.00  5.00]] at 0x7f8bfd8d9678)\n",
      "      - None \n",
      "    - None \n"
     ]
    }
   ],
   "source": [
    "def computational_graph(op, prefix=''):\n",
    "    print(f'{prefix}- {op} ' + (f'({op.variable.detach().numpy()} at {hex(id(op.variable))})' if op.__class__.__name__ == 'AccumulateGrad' else ''))\n",
    "    if op is not None: \n",
    "        for op in op.next_functions:\n",
    "            computational_graph(op[0], prefix+'  ')\n",
    "computational_graph(z.grad_fn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Index select (not needed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                                                              [    2     4     6     8     8]\n",
      "[ 1.00  1.00  1.00  1.00  2.00  6.00  0.00  1.00 -2.00] ----> [ 1.00  2.00  0.00 -2.00 -2.00]\n",
      "[ 0.00  0.00  3.00  0.00  3.00  0.00  0.00  0.00  6.00] <---- [ 3.00  3.00  0.00  3.00  3.00]\n"
     ]
    }
   ],
   "source": [
    "a = torch.tensor(\n",
    "    [1, 1, 1, 1, 2, 6, 0, 1, -2], dtype=torch.float, requires_grad=True)\n",
    "a_idx = torch.tensor(\n",
    "    [2, 4, 6, 8, 8])\n",
    "a_new = torch.index_select(a, index=a_idx, dim=0)\n",
    "\n",
    "grad_out = torch.full_like(a_new, 3) * (a_new != 0).float()\n",
    "a_new.backward(grad_out)\n",
    "\n",
    "print(' ' * 61, a_idx.numpy())\n",
    "print(a.detach().numpy(), '---->', a_new.detach().numpy())\n",
    "print(a.grad.numpy(), '<----', grad_out.detach().numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "class IndexSelectRelevance(torch.autograd.Function):\n",
    "    @staticmethod\n",
    "    def forward(ctx, src, dim, idx):\n",
    "        out = torch.index_select(src, dim, idx)\n",
    "        ctx.dim = dim\n",
    "        ctx.dim_size = src.shape[dim]\n",
    "        ctx.save_for_backward(src, idx, out)\n",
    "        return out\n",
    "\n",
    "    @staticmethod\n",
    "    def backward(ctx, rel_out):\n",
    "        src, idx, out = ctx.saved_tensors\n",
    "        return torch_scatter.scatter_add(rel_out, idx, dim=ctx.dim, dim_size=ctx.dim_size), None, None\n",
    "    \n",
    "def index_select(src, dim, index):\n",
    "    return IndexSelectRelevance.apply(src, dim, index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                                                              [    2     4     6     8     8]\n",
      "[ 1.00  1.00  1.00  1.00  2.00  6.00  0.00  1.00 -2.00] ----> [ 1.00  2.00  0.00 -2.00 -2.00]\n",
      "[ 0.00  0.00  3.00  0.00  3.00  0.00  0.00  0.00  6.00] <---- [ 3.00  3.00  0.00  3.00  3.00]\n"
     ]
    }
   ],
   "source": [
    "a = torch.tensor(\n",
    "    [1, 1, 1, 1, 2, 6, 0, 1, -2], dtype=torch.float, requires_grad=True)\n",
    "a_idx = torch.tensor(\n",
    "    [2, 4, 6, 8, 8])\n",
    "a_new = index_select(a, index=a_idx, dim=0)\n",
    "\n",
    "rel_out = torch.full_like(a_new, 3) * (a_new != 0).float()\n",
    "a_new.backward(rel_out)\n",
    "\n",
    "print(' ' * 61, a_idx.numpy())\n",
    "print(a.detach().numpy(), '---->', a_new.detach().numpy())\n",
    "print(a.grad.numpy(), '<----', rel_out.detach().numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Cat (not needed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 1.00  2.00  0.00] : [-1.00 -2.00  0.00] ----> [ 1.00  2.00  0.00 -1.00 -2.00  0.00]\n",
      "[ 3.00  3.00  0.00]   [ 3.00  3.00  0.00] <---- [ 3.00  3.00  0.00  3.00  3.00  0.00]\n"
     ]
    }
   ],
   "source": [
    "a = torch.tensor([ 1,  2, 0], dtype=torch.float, requires_grad=True)\n",
    "b = torch.tensor([-1, -2, 0], dtype=torch.float, requires_grad=True)\n",
    "c = torch.cat([a, b], dim=0)\n",
    "\n",
    "grad_out = torch.full_like(c, 3) * (c != 0).float()\n",
    "c.backward(grad_out)\n",
    "\n",
    "print(a.detach().numpy(), ':', b.detach().numpy(), '---->', c.detach().numpy())\n",
    "print(a.grad.numpy(), ' ', b.grad.numpy(), '<----', grad_out.detach().numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CatRelevance(torch.autograd.Function):\n",
    "    @staticmethod\n",
    "    def forward(ctx, dim, *tensors):\n",
    "        ctx.dim = dim\n",
    "        ctx.sizes = [t.shape[dim] for t in tensors]\n",
    "        return torch.cat(tensors, dim)\n",
    "\n",
    "    @staticmethod\n",
    "    def backward(ctx, rel_out):\n",
    "        return (None, *torch.split_with_sizes(rel_out, dim=ctx.dim, split_sizes=ctx.sizes))\n",
    "    \n",
    "def cat(tensors, dim=0):\n",
    "    return CatRelevance.apply(dim, *tensors)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 1.00  2.00  0.00] : [-1.00 -2.00  0.00] ----> [ 1.00  2.00  0.00 -1.00 -2.00  0.00]\n",
      "[ 3.00  3.00  0.00]   [ 3.00  3.00  0.00] <---- [ 3.00  3.00  0.00  3.00  3.00  0.00]\n"
     ]
    }
   ],
   "source": [
    "a = torch.tensor([ 1,  2, 0], dtype=torch.float, requires_grad=True)\n",
    "b = torch.tensor([-1, -2, 0], dtype=torch.float, requires_grad=True)\n",
    "c = cat((a, b), dim=0)\n",
    "\n",
    "rel_out = torch.full_like(c, 3) * (c != 0).float()\n",
    "c.backward(rel_out)\n",
    "\n",
    "print(a.detach().numpy(), ':', b.detach().numpy(), '---->', c.detach().numpy())\n",
    "print(a.grad.numpy(), ' ', b.grad.numpy(), '<----', rel_out.detach().numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "toc-hr-collapsed": false
   },
   "source": [
    "## Repeat tensor (not needed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[    3     0     2]\n",
      "[ 7.00  2.00  0.00] ----> [ 7.00  7.00  7.00  0.00  0.00]\n",
      "[ 3.00  0.00  0.00] <---- [ 1.00  1.00  1.00  0.00  0.00]\n"
     ]
    }
   ],
   "source": [
    "a = torch.tensor([ 7, 2, 0], dtype=torch.float, requires_grad=True)\n",
    "a_repeats = torch.tensor([ 3, 0, 2])\n",
    "a_new = tg.utils.repeat_tensor(a, a_repeats, dim=0)\n",
    "\n",
    "grad_out = torch.ones_like(a_new) * (a_new != 0).float()\n",
    "a_new.backward(grad_out)\n",
    "\n",
    "print(a_repeats.numpy())\n",
    "print(a.detach().numpy(), '---->', a_new.detach().numpy())\n",
    "print(a.grad.numpy(), '<----', grad_out.numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "def repeat_tensor(src, repeats, dim=0):\n",
    "    idx = src.new_tensor(np.arange(len(repeats)).repeat(repeats.cpu().numpy()), dtype=torch.long)\n",
    "    return torch.index_select(src, dim, idx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[    3     0     2]\n",
      "[ 7.00  2.00  0.00] ----> [ 7.00  7.00  7.00  0.00  0.00]\n",
      "[ 3.00  0.00  0.00] <---- [ 1.00  1.00  1.00  0.00  0.00]\n"
     ]
    }
   ],
   "source": [
    "a = torch.tensor([ 7, 2, 0], dtype=torch.float, requires_grad=True)\n",
    "a_repeats = torch.tensor([ 3, 0, 2])\n",
    "a_new = repeat_tensor(a, a_repeats, dim=0)\n",
    "\n",
    "rel_out = torch.ones_like(a_new) * (a_new != 0).float()\n",
    "a_new.backward(rel_out)\n",
    "\n",
    "print(a_repeats.numpy())\n",
    "print(a.detach().numpy(), '---->', a_new.detach().numpy())\n",
    "print(a.grad.numpy(), '<----', rel_out.numpy())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:gn-exp]",
   "language": "python",
   "name": "conda-env-gn-exp-py"
  },
  "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"
  },
  "toc-autonumbering": false,
  "toc-showmarkdowntxt": false,
  "toc-showtags": false
 },
 "nbformat": 4,
 "nbformat_minor": 2
}