{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 383,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.autograd as autograd\n",
    "\n",
    "from torch import Tensor\n",
    "from typing import *\n",
    "\n",
    "from torch_scatter import scatter\n",
    "from torch_geometric.utils import softmax\n",
    "\n",
    "import triton\n",
    "import triton.language as tl\n",
    "\n",
    "from copy import deepcopy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 319,
   "metadata": {},
   "outputs": [],
   "source": [
    "# class SoftmaxSumFunction(autograd.Function):\n",
    "#     @staticmethod\n",
    "#     def forward(\n",
    "#         ctx: autograd.function.FunctionCtx,\n",
    "#         x: Tensor,\n",
    "#         a: Tensor,\n",
    "#         _x: Tensor,\n",
    "#         index: Tensor,\n",
    "#         mask: Tensor,\n",
    "#         exp_a: Tensor,\n",
    "#         sum_exp_a: Tensor,\n",
    "#         sum_exp_ah: Tensor,\n",
    "#         training: bool,\n",
    "#     ):\n",
    "#         dlt_exp_ah = exp_a[mask, None] * (x - _x)\n",
    "#         sum_exp_ah = scatter(dlt_exp_ah, index, dim=0, out=sum_exp_ah.clone(), reduce=\"sum\")\n",
    "#         out = sum_exp_ah / sum_exp_a[:, None]\n",
    "#         if training:\n",
    "#             w = exp_a[mask] / sum_exp_a.index_select(0, index)\n",
    "#             ctx.save_for_backward(out, x, w, index)\n",
    "#         return out\n",
    "    \n",
    "#     @staticmethod\n",
    "#     def backward(\n",
    "#         ctx: autograd.function.FunctionCtx,\n",
    "#         grad: Tensor,\n",
    "#     ):\n",
    "#         out, x, w, index = ctx.saved_tensors\n",
    "#         g = grad.index_select(0, index)\n",
    "#         o =  out.index_select(0, index)\n",
    "\n",
    "#         grad_x = g * w[:, None]\n",
    "#         grad_a = (g * (x - o)).sum(dim=-1) * w\n",
    "#         return grad_x, grad_a, None, None, None, None, None, None, None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 320,
   "metadata": {},
   "outputs": [],
   "source": [
    "# class SoftmaxSum(nn.Module):\n",
    "#     def __init__(self,\n",
    "#         num_nodes: int,\n",
    "#         num_edges: int,\n",
    "#         num_features: int,\n",
    "#         dtype: Optional[torch.dtype] = None,\n",
    "#     ) -> None:\n",
    "#         super().__init__()\n",
    "#         self.num_nodes = num_nodes\n",
    "#         self.num_edges = num_edges\n",
    "\n",
    "#         self.register_buffer(\"max_hist_a\", torch.zeros(num_nodes, dtype=dtype))\n",
    "#         self.register_buffer(\"exp_hist_a\", torch.zeros(num_edges, dtype=dtype))\n",
    "#         self.register_buffer(\"sum_exp_a\", torch.zeros(num_nodes, dtype=dtype))\n",
    "#         self.register_buffer(\"sum_exp_ah\", torch.zeros(num_nodes, num_features, dtype=dtype))\n",
    "    \n",
    "#     @torch.no_grad()\n",
    "#     def update_snapshot(self, x: Tensor, a: Tensor, index: Tensor):\n",
    "#         max_a = scatter(a, index, dim=0, dim_size=self.num_nodes, reduce=\"max\")\n",
    "#         exp_a = (a - max_a.index_select(0, index)).exp_()\n",
    "#         self.get_buffer(\"max_hist_a\")[:] = max_a\n",
    "#         self.get_buffer(\"exp_hist_a\")[:] = exp_a\n",
    "\n",
    "#         sum_exp_a = scatter(exp_a, index, dim=0, dim_size=self.num_nodes, reduce=\"sum\") + 1e-10\n",
    "#         sum_exp_ah = scatter(exp_a[:,None] * x, index, dim=0, dim_size=self.num_nodes, reduce=\"sum\")\n",
    "#         self.get_buffer(\"sum_exp_a\")[:] = sum_exp_a\n",
    "#         self.get_buffer(\"sum_exp_ah\")[:] = sum_exp_ah\n",
    "#         return self\n",
    "    \n",
    "#     @torch.no_grad()\n",
    "#     def update_attention(self, x: Tensor, a: Tensor, _x: Tensor, index: Tensor, mask: Tensor):\n",
    "#         exp_a = (a - self.get_buffer(\"max_hist_a\").index_select(0, index)).exp_()\n",
    "#         exp_ah = exp_a[:, None] * x\n",
    "\n",
    "#         exp_hist_a = self.get_buffer(\"exp_hist_a\")[mask]\n",
    "#         exp_hist_ah = exp_hist_a[:, None] * _x\n",
    "\n",
    "#         self.get_buffer(\"exp_hist_a\")[mask] = exp_a\n",
    "#         scatter(exp_a - exp_hist_a, index, dim=0, out=self.get_buffer(\"sum_exp_a\"), reduce=\"sum\")\n",
    "#         scatter(exp_ah - exp_hist_ah, index, dim=0, out=self.get_buffer(\"sum_exp_ah\"), reduce=\"sum\")\n",
    "#         return self\n",
    "    \n",
    "#     def forward(self, x: Tensor, a: Tensor, _x: Tensor, index: Tensor, mask: Tensor):\n",
    "#         return SoftmaxSumFunction.apply(x, a, _x, index, mask,\n",
    "#                                  self.get_buffer(\"exp_hist_a\"),\n",
    "#                                  self.get_buffer(\"sum_exp_a\"),\n",
    "#                                  self.get_buffer(\"sum_exp_ah\"),\n",
    "#                                  self.training)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 321,
   "metadata": {},
   "outputs": [],
   "source": [
    "# num_nodes = 200\n",
    "# num_edges = 100\n",
    "# num_features = 8\n",
    "# for i in range(10):\n",
    "#     x = torch.randn(num_edges, num_features)\n",
    "#     a = torch.randn(num_edges)\n",
    "#     index = torch.randint(num_nodes, size=(num_edges,), dtype=torch.long)\n",
    "#     mask = torch.ones(num_edges, dtype=torch.bool)\n",
    "\n",
    "#     ss = SoftmaxSum(num_nodes, num_edges, num_features)\n",
    "#     ss.update_snapshot(x, a, index)\n",
    "#     ss.update_attention(x, a, x, index, mask)\n",
    "\n",
    "#     x.requires_grad_()\n",
    "#     a.requires_grad_()\n",
    "#     ss(x, a, x, index, mask).sum().backward()\n",
    "#     grad_x0 = x.grad.clone()\n",
    "#     grad_a0 = a.grad.clone()\n",
    "\n",
    "#     x.grad = None\n",
    "#     a.grad = None\n",
    "#     t = softmax(a, index, num_nodes=num_nodes)[:,None] * x\n",
    "#     scatter(t, index, dim=0, dim_size=num_nodes, reduce=\"sum\").sum().backward()\n",
    "\n",
    "#     # print(torch.norm(grad_x0 - x.grad), torch.norm(grad_a0 - a.grad))\n",
    "#     # print(torch.abs(grad_x0 - x.grad).max(), torch.abs(grad_a0 - a.grad).max())\n",
    "#     # print(grad_x0)\n",
    "#     # print(grad_a0)\n",
    "#     print(grad_a0 - a.grad)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 355,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SoftmaxSumFunction(autograd.Function):\n",
    "    @staticmethod\n",
    "    def forward(\n",
    "        ctx: autograd.function.FunctionCtx,\n",
    "        x1: Tensor,\n",
    "        a1: Tensor,\n",
    "        x0: Tensor,\n",
    "        a0: Tensor,\n",
    "        index: Tensor,\n",
    "        max_a: Tensor,\n",
    "        sum_exp_a: Tensor,\n",
    "        sum_exp_ah: Tensor,\n",
    "        training: bool,\n",
    "    ):\n",
    "        exp_a = (a0 - max_a.index_select(0, index)).exp_()\n",
    "        dlt_exp_ah = exp_a[:, :, None] * (x1 - x0)\n",
    "        sum_exp_ah = scatter(dlt_exp_ah, index, dim=0, out=sum_exp_ah.clone(), reduce=\"sum\")\n",
    "        out = sum_exp_ah / sum_exp_a[:, :, None]\n",
    "        if training:\n",
    "            w = exp_a / sum_exp_a.index_select(0, index)\n",
    "            ctx.save_for_backward(out, x1, w, index)\n",
    "        return out\n",
    "    \n",
    "    @staticmethod\n",
    "    def backward(\n",
    "        ctx: autograd.function.FunctionCtx,\n",
    "        grad: Tensor,\n",
    "    ):\n",
    "        out, x, w, index = ctx.saved_tensors\n",
    "        g = grad.index_select(0, index)\n",
    "        o =  out.index_select(0, index)\n",
    "\n",
    "        grad_x = g * w[:, :, None]\n",
    "        grad_a = (g * (x - o)).sum(dim=-1) * w\n",
    "        return grad_x, grad_a, None, None, None, None, None, None, None\n",
    "\n",
    "class SoftmaxSum(nn.Module):\n",
    "    def __init__(self,\n",
    "        num_nodes: int,\n",
    "        num_edges: int,\n",
    "        num_heads: int,\n",
    "        num_features: int,\n",
    "        dtype: Optional[torch.dtype] = None,\n",
    "    ) -> None:\n",
    "        super().__init__()\n",
    "        self.num_nodes = num_nodes\n",
    "        self.num_edges = num_edges\n",
    "        self.num_heads = num_heads\n",
    "\n",
    "        self.register_buffer(\"max_a\", torch.zeros(num_nodes, num_heads, dtype=dtype))\n",
    "        self.register_buffer(\"sum_exp_a\", torch.zeros(num_nodes, num_heads, dtype=dtype))\n",
    "        self.register_buffer(\"sum_exp_ah\", torch.zeros(num_nodes, num_heads, num_features, dtype=dtype))\n",
    "    \n",
    "    @torch.no_grad()\n",
    "    def update_snapshot(self, x: Tensor, a: Tensor, index: Tensor):\n",
    "        max_a = scatter(a, index, dim=0, dim_size=self.num_nodes, reduce=\"max\")\n",
    "        exp_a = (a - max_a.index_select(0, index)).exp_()\n",
    "        self.get_buffer(\"max_a\")[:] = max_a\n",
    "\n",
    "        sum_exp_a = scatter(exp_a, index, dim=0, dim_size=self.num_nodes, reduce=\"sum\") + 1e-10\n",
    "        sum_exp_ah = scatter(exp_a[:, :, None] * x, index, dim=0, dim_size=self.num_nodes, reduce=\"sum\")\n",
    "        self.get_buffer(\"sum_exp_a\")[:] = sum_exp_a\n",
    "        self.get_buffer(\"sum_exp_ah\")[:] = sum_exp_ah\n",
    "        return self\n",
    "    \n",
    "    @torch.no_grad()\n",
    "    def update_attention(self, x1: Tensor, a1: Tensor, x0: Tensor, a0: Tensor, index: Tensor):\n",
    "        max_hist_a = self.get_buffer(\"max_a\").index_select(0, index)\n",
    "        exp_a0 = (a0 - max_hist_a).exp_()\n",
    "        exp_a1 = (a1 - max_hist_a).exp_()\n",
    "        scatter(exp_a1 - exp_a0, index, dim=0, out=self.get_buffer(\"sum_exp_a\"), reduce=\"sum\")\n",
    "\n",
    "        exp_ah0 = exp_a0[:, :, None] * x0\n",
    "        exp_ah1 = exp_a1[:, :, None] * x1\n",
    "        scatter(exp_ah1 - exp_ah0, index, dim=0, out=self.get_buffer(\"sum_exp_ah\"), reduce=\"sum\")\n",
    "        return self\n",
    "    \n",
    "    def forward(self, x1: Tensor, a1: Tensor, x0: Tensor, a0: Tensor, index: Tensor) -> Tensor:\n",
    "        return SoftmaxSumFunction.apply(x1, a1, x0, a0, index,\n",
    "                                 self.get_buffer(\"max_a\"),\n",
    "                                 self.get_buffer(\"sum_exp_a\"),\n",
    "                                 self.get_buffer(\"sum_exp_ah\"),\n",
    "                                 self.training)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 359,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.) tensor(6.2413e-06)\n",
      "tensor([[-2.1367, -3.5242],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.8726, -6.0982]])\n",
      "tensor([[-2.1367, -3.5242],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.8726, -6.0982]])\n",
      "tensor(0.) tensor(7.5973e-06)\n",
      "tensor([[ 4.1189, -1.3866],\n",
      "        [ 6.2586, -4.8426],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 2.7955,  8.4197],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000]])\n",
      "tensor([[ 4.1189, -1.3866],\n",
      "        [ 6.2586, -4.8426],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 2.7955,  8.4197],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000]])\n",
      "tensor(0.) tensor(6.0887e-06)\n",
      "tensor([[ 0.0000,  0.0000],\n",
      "        [-1.7955,  4.7750],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [-3.7791, -3.4205],\n",
      "        [ 3.1130,  3.6583],\n",
      "        [ 0.0000,  0.0000]])\n",
      "tensor([[ 0.0000,  0.0000],\n",
      "        [-1.7955,  4.7750],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [-3.7791, -3.4205],\n",
      "        [ 3.1130,  3.6583],\n",
      "        [ 0.0000,  0.0000]])\n",
      "tensor(0.) tensor(5.9715e-06)\n",
      "tensor([[ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 9.4426, -2.2261],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 2.7483,  3.0900],\n",
      "        [-9.4426,  2.2261]])\n",
      "tensor([[ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 9.4426, -2.2261],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 2.7483,  3.0900],\n",
      "        [-9.4426,  2.2261]])\n",
      "tensor(0.) tensor(7.1696e-06)\n",
      "tensor([[ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 2.2263,  0.6597],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [-2.2893, -4.1325],\n",
      "        [-3.9234, -3.9451],\n",
      "        [ 0.0000,  0.0000]])\n",
      "tensor([[ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 2.2263,  0.6597],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [-2.2893, -4.1325],\n",
      "        [-3.9234, -3.9451],\n",
      "        [ 0.0000,  0.0000]])\n",
      "tensor(0.) tensor(6.6798e-06)\n",
      "tensor([[ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 1.8056,  3.9442],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [-2.4050, -5.5951]])\n",
      "tensor([[ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 1.8056,  3.9442],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [-2.4050, -5.5951]])\n",
      "tensor(0.) tensor(4.7071e-06)\n",
      "tensor([[  0.0000,   0.0000],\n",
      "        [ -3.6159, -10.4370],\n",
      "        [  0.0000,   0.0000],\n",
      "        [  0.0000,   0.0000],\n",
      "        [  0.0000,   0.0000],\n",
      "        [  0.0000,   0.0000],\n",
      "        [  2.1697,   4.4867],\n",
      "        [  0.8153,   9.0103]])\n",
      "tensor([[  0.0000,   0.0000],\n",
      "        [ -3.6159, -10.4370],\n",
      "        [  0.0000,   0.0000],\n",
      "        [  0.0000,   0.0000],\n",
      "        [  0.0000,   0.0000],\n",
      "        [  0.0000,   0.0000],\n",
      "        [  2.1697,   4.4867],\n",
      "        [  0.8153,   9.0103]])\n",
      "tensor(0.) tensor(6.2858e-06)\n",
      "tensor([[ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 5.5491, -3.5294],\n",
      "        [ 2.2001, -5.5199],\n",
      "        [ 1.6753,  1.5286],\n",
      "        [ 5.7522, -2.5286]])\n",
      "tensor([[ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 5.5491, -3.5294],\n",
      "        [ 2.2001, -5.5199],\n",
      "        [ 1.6753,  1.5286],\n",
      "        [ 5.7522, -2.5286]])\n",
      "tensor(0.) tensor(5.8452e-06)\n",
      "tensor([[ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [-0.3675,  3.0714],\n",
      "        [ 0.4008, -9.1039],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000]])\n",
      "tensor([[ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [-0.3675,  3.0714],\n",
      "        [ 0.4008, -9.1039],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000]])\n",
      "tensor(0.) tensor(6.5390e-06)\n",
      "tensor([[ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 4.4908,  0.8768],\n",
      "        [ 0.1810,  2.1797],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [-4.4908, -0.8768]])\n",
      "tensor([[ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 4.4908,  0.8768],\n",
      "        [ 0.1810,  2.1797],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000],\n",
      "        [-4.4908, -0.8768]])\n"
     ]
    }
   ],
   "source": [
    "num_nodes = 200\n",
    "num_edges = 100\n",
    "num_heads = 2\n",
    "num_features = 256\n",
    "for i in range(10):\n",
    "    x = torch.randn(num_edges, num_heads, num_features)\n",
    "    a = torch.randn(num_edges, num_heads)\n",
    "    index = torch.randint(num_nodes, size=(num_edges,), dtype=torch.long)\n",
    "\n",
    "    ss = SoftmaxSum(num_nodes, num_edges, num_heads, num_features)\n",
    "    ss.update_snapshot(x, a, index)\n",
    "    ss.update_attention(x, a, x, a, index)\n",
    "\n",
    "    x.requires_grad_()\n",
    "    a.requires_grad_()\n",
    "    ss(x, a, x, a, index).sum().backward()\n",
    "    grad_x0 = x.grad.clone()\n",
    "    grad_a0 = a.grad.clone()\n",
    "\n",
    "    x.grad = None\n",
    "    a.grad = None\n",
    "    t = softmax(a, index, num_nodes=num_nodes)[:, :, None] * x\n",
    "    scatter(t, index, dim=0, dim_size=num_nodes, reduce=\"sum\").sum().backward()\n",
    "\n",
    "    print(torch.norm(grad_x0 - x.grad), torch.norm(grad_a0 - a.grad))\n",
    "    # print(torch.abs(grad_x0 - x.grad).max(), torch.abs(grad_a0 - a.grad).max())\n",
    "    # print(grad_x0)\n",
    "    print(grad_a0[:8, :])\n",
    "    print(a.grad[:8, :])\n",
    "    # print(grad_a0 - a.grad)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 399,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GATConv(nn.Module):\n",
    "    def __init__(self,\n",
    "        in_channels: int,\n",
    "        out_channels: int,\n",
    "        heads: int = 1,\n",
    "        concat: bool = False,\n",
    "        negative_slope: float = 0.2,\n",
    "        bias: bool = True,\n",
    "        **kwargs\n",
    "    ) -> None:\n",
    "        super().__init__()\n",
    "        \n",
    "        self.in_channels = in_channels\n",
    "        self.out_channels = out_channels\n",
    "        self.heads = heads\n",
    "        self.concat = concat\n",
    "        self.negative_slope = negative_slope\n",
    "        self._softmax_sum_list = [None]\n",
    "        \n",
    "        self.weight = nn.Parameter(torch.Tensor(in_channels, heads * out_channels))\n",
    "        self.att_src = nn.Parameter(torch.Tensor(1, heads, out_channels))\n",
    "        self.att_dst = nn.Parameter(torch.Tensor(1, heads, out_channels))\n",
    "        \n",
    "        if bias and concat:\n",
    "            self.bias = nn.Parameter(torch.Tensor(heads * out_channels))\n",
    "        elif bias and not concat:\n",
    "            self.bias = nn.Parameter(torch.Tensor(out_channels))\n",
    "        else:\n",
    "            self.bias = None\n",
    "\n",
    "        self.reset_parameters()\n",
    "    \n",
    "    def reset_parameters(self):\n",
    "        nn.init.xavier_normal_(self.weight)\n",
    "        nn.init.xavier_normal_(self.att_src)\n",
    "        nn.init.xavier_normal_(self.att_dst)\n",
    "        if self.bias is not None:\n",
    "            nn.init.zeros_(self.bias)\n",
    "    \n",
    "    @property\n",
    "    def softmax_sum(self) -> Optional[SoftmaxSum]:\n",
    "        return self._softmax_sum_list[0]\n",
    "    \n",
    "    def set_softmax_sum(self, m: Optional[SoftmaxSum]):\n",
    "        self._softmax_sum_list = [m]\n",
    "        return self\n",
    "    \n",
    "    def train(self, mode: bool = True):\n",
    "        if self.softmax_sum is not None:\n",
    "            self.softmax_sum.train(mode)\n",
    "        return super().train(mode)\n",
    "    \n",
    "    def update_softmax_snapshot(self,\n",
    "        x: Tensor,\n",
    "        edge_index: Tensor,\n",
    "        snap_net: nn.Module,\n",
    "    ):\n",
    "        snap_net: GATConv = snap_net.eval()\n",
    "        with torch.no_grad():\n",
    "            x, a = snap_net._att_x(x, edge_index)\n",
    "        self.softmax_sum.update_snapshot(x, a, edge_index[1])\n",
    "\n",
    "    def _att_x(self, x: Tensor, edge_index: Tensor) -> Tuple[Tensor, Tensor]:\n",
    "        H, C = self.heads, self.out_channels\n",
    "        x = (x @ self.weight).view(-1, H, C)\n",
    "\n",
    "        src_a = (x * self.att_src).sum(dim=-1)\n",
    "        src_a = src_a[edge_index[0]]\n",
    "        dst_a = (x * self.att_dst).sum(dim=-1)\n",
    "        dst_a = dst_a[edge_index[1]]\n",
    "\n",
    "        x = x[edge_index[0]]\n",
    "        a = F.leaky_relu(src_a + dst_a, self.negative_slope)\n",
    "        return x, a\n",
    "    \n",
    "    def forward(self,\n",
    "        x1: Tensor,\n",
    "        x0: Tensor,\n",
    "        edge_index: Tensor,\n",
    "        snap_net: nn.Module,\n",
    "        update_att: bool,\n",
    "    ) -> Tensor:\n",
    "        H, C = self.heads, self.out_channels\n",
    "        snap_net: GATConv = snap_net.eval()\n",
    "\n",
    "        index = edge_index[1]\n",
    "        with torch.no_grad():\n",
    "            x0, a0 = snap_net._att_x(x0, edge_index)\n",
    "            if update_att:\n",
    "                x, a = snap_net._att_x(x1, edge_index)\n",
    "                self.softmax_sum.update_attention(x, a, x0, a0, index)\n",
    "                x0, a0 = x, a\n",
    "        \n",
    "        x1, a1 = self._att_x(x1, edge_index)\n",
    "        out = self.softmax_sum.forward(x1, a1, x0, a0, index)\n",
    "        \n",
    "        if self.concat:\n",
    "            out = out.view(-1, H * C)\n",
    "        else:\n",
    "            out = out.mean(dim=1)\n",
    "\n",
    "        if self.bias is not None:\n",
    "            out += self.bias\n",
    "        return out\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 400,
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def slim_local_graph(src_mask: Tensor, edge_index: Tensor) -> Tuple[Tensor, Tensor, Tensor]:\n",
    "    edge_mask = src_mask[edge_index[0]]\n",
    "    slim_edge_index = edge_index[:, edge_mask]\n",
    "    \n",
    "    node_mask = src_mask.clone().index_fill_(0, slim_edge_index[1], 1)\n",
    "\n",
    "    idx = torch.where(node_mask)[0]\n",
    "    imp = torch.empty_like(src_mask, dtype=torch.long).fill_((2**62-1)*2+1)\n",
    "    imp[idx] = torch.arange(idx.size(0))\n",
    "    \n",
    "    slim_edge_index = imp[slim_edge_index.flatten()].view_as(slim_edge_index)\n",
    "    return node_mask, edge_mask, slim_edge_index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 401,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([200])\n",
      "torch.Size([100])\n",
      "torch.Size([2, 35])\n"
     ]
    }
   ],
   "source": [
    "src_mask = torch.rand(num_nodes) < 0.3\n",
    "edge_index = torch.randint(num_nodes, size=(2, num_edges), dtype=torch.long)\n",
    "node_mask, edge_mask, slim_edge_index = slim_local_graph(src_mask, edge_index)\n",
    "print(node_mask.size())\n",
    "print(edge_mask.size())\n",
    "print(slim_edge_index.size())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 484,
   "metadata": {},
   "outputs": [],
   "source": [
    "net = GATConv(3, num_features, num_heads)\n",
    "att = SoftmaxSum(num_nodes, num_edges, num_heads, num_features)\n",
    "net.set_softmax_sum(att)\n",
    "\n",
    "snap_net = deepcopy(net)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 485,
   "metadata": {},
   "outputs": [],
   "source": [
    "x1 = torch.randn(num_nodes, 3)\n",
    "x0 = torch.randn(num_nodes, 3)\n",
    "edge_index = torch.randint(num_nodes, size=(2, num_edges), dtype=torch.long)\n",
    "\n",
    "net.update_softmax_snapshot(x0, edge_index, snap_net)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 511,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.1482, -0.0744, -0.0213,  ..., -0.0421, -0.1849, -0.1039],\n",
       "        [-0.0075,  0.0321, -0.0046,  ...,  0.0154,  0.0506, -0.0110],\n",
       "        [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],\n",
       "        ...,\n",
       "        [ 0.0878, -0.0887, -0.0550,  ..., -0.0982, -0.1763, -0.0175],\n",
       "        [ 0.1193,  0.0695, -0.0637,  ...,  0.0060,  0.0391, -0.1508],\n",
       "        [ 0.0294,  0.0186,  0.0140,  ...,  0.0327,  0.0128, -0.0487]],\n",
       "       grad_fn=<AddBackward0>)"
      ]
     },
     "execution_count": 511,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = net.forward(x1, x0, edge_index, snap_net, update_att=True)\n",
    "x0 = x1\n",
    "x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mpi",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
