test_gat_layer_implementation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: appnope==0.1.3 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 1)) (0.1.3)\n",
      "Requirement already satisfied: asttokens==2.2.1 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 2)) (2.2.1)\n",
      "Requirement already satisfied: attrs==22.1.0 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 3)) (22.1.0)\n",
      "Requirement already satisfied: backcall==0.2.0 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 4)) (0.2.0)\n",
      "Requirement already satisfied: backports.zoneinfo==0.2.1 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 5)) (0.2.1)\n",
      "Requirement already satisfied: beautifulsoup4==4.11.1 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 6)) (4.11.1)\n",
      "Requirement already satisfied: bleach==5.0.1 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 7)) (5.0.1)\n",
      "Requirement already satisfied: cachetools==5.2.0 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 8)) (5.2.0)\n",
      "Requirement already satisfied: certifi==2022.12.7 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 9)) (2022.12.7)\n",
      "Requirement already satisfied: chardet==3.0.4 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 10)) (3.0.4)\n",
      "Requirement already satisfied: charset-normalizer==2.1.1 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 11)) (2.1.1)\n",
      "Requirement already satisfied: comm==0.1.2 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 12)) (0.1.2)\n",
      "Requirement already satisfied: contourpy==1.0.6 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 13)) (1.0.6)\n",
      "Requirement already satisfied: cycler==0.11.0 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 14)) (0.11.0)\n",
      "Requirement already satisfied: debugpy==1.6.4 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 15)) (1.6.4)\n",
      "Requirement already satisfied: decorator==5.1.1 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 16)) (5.1.1)\n",
      "Requirement already satisfied: defusedxml==0.7.1 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 17)) (0.7.1)\n",
      "Requirement already satisfied: docopt==0.6.2 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 18)) (0.6.2)\n",
      "Requirement already satisfied: entrypoints==0.4 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 19)) (0.4)\n",
      "Requirement already satisfied: executing==1.2.0 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 20)) (1.2.0)\n",
      "Requirement already satisfied: fastjsonschema==2.16.2 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 21)) (2.16.2)\n",
      "Requirement already satisfied: fonttools==4.38.0 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 22)) (4.38.0)\n",
      "Requirement already satisfied: google-auth==1.4.2 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 23)) (1.4.2)\n",
      "Requirement already satisfied: idna==3.4 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 24)) (3.4)\n",
      "Requirement already satisfied: importlib-metadata==5.1.0 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 25)) (5.1.0)\n",
      "Requirement already satisfied: importlib-resources==5.10.1 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 26)) (5.10.1)\n",
      "Requirement already satisfied: ipykernel==6.19.2 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 27)) (6.19.2)\n",
      "Requirement already satisfied: ipython==8.7.0 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 28)) (8.7.0)\n",
      "Requirement already satisfied: ipython-genutils==0.2.0 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 29)) (0.2.0)\n",
      "Requirement already satisfied: isort==5.11.2 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 30)) (5.11.2)\n",
      "Requirement already satisfied: jedi==0.18.2 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 31)) (0.18.2)\n",
      "Requirement already satisfied: Jinja2==3.1.2 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 32)) (3.1.2)\n",
      "Requirement already satisfied: joblib==1.2.0 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 33)) (1.2.0)\n",
      "Requirement already satisfied: Js2Py==0.74 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 34)) (0.74)\n",
      "Requirement already satisfied: jupyter_client==7.4.8 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 35)) (7.4.8)\n",
      "Requirement already satisfied: jupyter_core==5.1.0 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 36)) (5.1.0)\n",
      "Requirement already satisfied: jupyterlab-pygments==0.2.2 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 37)) (0.2.2)\n",
      "Requirement already satisfied: kiwisolver==1.4.4 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 38)) (1.4.4)\n",
      "Requirement already satisfied: MarkupSafe==2.1.1 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 39)) (2.1.1)\n",
      "Requirement already satisfied: matplotlib==3.6.2 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 40)) (3.6.2)\n",
      "Requirement already satisfied: matplotlib-inline==0.1.6 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 41)) (0.1.6)\n",
      "Requirement already satisfied: mistune==2.0.4 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 42)) (2.0.4)\n",
      "Requirement already satisfied: nest-asyncio==1.5.6 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 43)) (1.5.6)\n",
      "Requirement already satisfied: networkx==2.8.8 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 44)) (2.8.8)\n",
      "Requirement already satisfied: numpy==1.23.5 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 45)) (1.23.5)\n",
      "Requirement already satisfied: packaging==22.0 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 46)) (22.0)\n",
      "Requirement already satisfied: pandas==1.5.2 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 47)) (1.5.2)\n",
      "Requirement already satisfied: pandocfilters==1.5.0 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 48)) (1.5.0)\n",
      "Requirement already satisfied: parso==0.8.3 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 49)) (0.8.3)\n",
      "Requirement already satisfied: pexpect==4.8.0 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 50)) (4.8.0)\n",
      "Requirement already satisfied: pickleshare==0.7.5 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 51)) (0.7.5)\n",
      "Requirement already satisfied: Pillow==9.3.0 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 52)) (9.3.0)\n",
      "Requirement already satisfied: pkgutil_resolve_name==1.3.10 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 53)) (1.3.10)\n",
      "Requirement already satisfied: platformdirs==2.6.0 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 54)) (2.6.0)\n",
      "Requirement already satisfied: portpicker==1.2.0 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 55)) (1.2.0)\n",
      "Requirement already satisfied: prompt-toolkit==3.0.36 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 56)) (3.0.36)\n",
      "Requirement already satisfied: psutil==5.9.4 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 57)) (5.9.4)\n",
      "Requirement already satisfied: ptyprocess==0.7.0 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 58)) (0.7.0)\n",
      "Requirement already satisfied: pure-eval==0.2.2 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 59)) (0.2.2)\n",
      "Requirement already satisfied: pyasn1==0.4.8 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 60)) (0.4.8)\n",
      "Requirement already satisfied: pyasn1-modules==0.2.8 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 61)) (0.2.8)\n",
      "Requirement already satisfied: pycodestyle==2.10.0 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 62)) (2.10.0)\n",
      "Requirement already satisfied: pyg-lib==0.1.0+pt113 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 63)) (0.1.0+pt113)\n",
      "Requirement already satisfied: Pygments==2.13.0 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 64)) (2.13.0)\n",
      "Requirement already satisfied: pyjsparser==2.7.1 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 65)) (2.7.1)\n",
      "Requirement already satisfied: pyparsing==3.0.9 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 66)) (3.0.9)\n",
      "Requirement already satisfied: PyPrind==2.11.3 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 67)) (2.11.3)\n",
      "Requirement already satisfied: pyrsistent==0.19.2 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 68)) (0.19.2)\n",
      "Requirement already satisfied: pySmartDL==1.3.4 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 69)) (1.3.4)\n",
      "Requirement already satisfied: python-dateutil==2.8.2 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 70)) (2.8.2)\n",
      "Requirement already satisfied: pytz==2022.6 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 71)) (2022.6)\n",
      "Requirement already satisfied: pytz-deprecation-shim==0.1.0.post0 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 72)) (0.1.0.post0)\n",
      "Requirement already satisfied: pyzmq==24.0.1 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 73)) (24.0.1)\n",
      "Requirement already satisfied: requests==2.28.1 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 74)) (2.28.1)\n",
      "Requirement already satisfied: rsa==4.9 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 75)) (4.9)\n",
      "Requirement already satisfied: scikit-learn==1.2.0 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 76)) (1.2.0)\n",
      "Requirement already satisfied: scipy==1.9.3 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 77)) (1.9.3)\n",
      "Requirement already satisfied: seaborn==0.12.1 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 78)) (0.12.1)\n",
      "Requirement already satisfied: simplegeneric==0.8.1 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 79)) (0.8.1)\n",
      "Requirement already satisfied: six==1.16.0 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 80)) (1.16.0)\n",
      "Requirement already satisfied: sklearn==0.0.post1 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 81)) (0.0.post1)\n",
      "Requirement already satisfied: soupsieve==2.3.2.post1 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 82)) (2.3.2.post1)\n",
      "Requirement already satisfied: stack-data==0.6.2 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 83)) (0.6.2)\n",
      "Requirement already satisfied: terminado==0.13.3 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 84)) (0.13.3)\n",
      "Requirement already satisfied: threadpoolctl==3.1.0 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 85)) (3.1.0)\n",
      "Requirement already satisfied: tinycss2==1.2.1 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 86)) (1.2.1)\n",
      "Requirement already satisfied: tomli==2.0.1 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 87)) (2.0.1)\n",
      "Requirement already satisfied: torch==1.13.0 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 88)) (1.13.0)\n",
      "Requirement already satisfied: torch-cluster==1.6.0 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 89)) (1.6.0)\n",
      "Requirement already satisfied: torch-geometric==2.2.0 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 90)) (2.2.0)\n",
      "Requirement already satisfied: torch-scatter==2.1.0 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 91)) (2.1.0)\n",
      "Requirement already satisfied: torch-sparse==0.6.15 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 92)) (0.6.15)\n",
      "Requirement already satisfied: torch-spline-conv==1.2.1 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 93)) (1.2.1)\n",
      "Requirement already satisfied: torchmetrics==0.11.0 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 94)) (0.11.0)\n",
      "Requirement already satisfied: tornado==6.2 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 95)) (6.2)\n",
      "Requirement already satisfied: tqdm==4.64.1 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 96)) (4.64.1)\n",
      "Requirement already satisfied: traitlets==5.7.1 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 97)) (5.7.1)\n",
      "Requirement already satisfied: typing_extensions==4.4.0 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 98)) (4.4.0)\n",
      "Requirement already satisfied: tzdata==2022.7 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 99)) (2022.7)\n",
      "Requirement already satisfied: tzlocal==4.2 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 100)) (4.2)\n",
      "Requirement already satisfied: urllib3==1.26.13 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 101)) (1.26.13)\n",
      "Requirement already satisfied: wcwidth==0.2.5 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 102)) (0.2.5)\n",
      "Requirement already satisfied: webencodings==0.5.1 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 103)) (0.5.1)\n",
      "Requirement already satisfied: yapf==0.32.0 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 104)) (0.32.0)\n",
      "Requirement already satisfied: zipp==3.11.0 in ./grl-env/lib/python3.8/site-packages (from -r requirements.txt (line 105)) (3.11.0)\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip install -r requirements.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch_geometric.nn as gnn\n",
    "import random\n",
    "from torch_geometric.utils import softmax, scatter\n",
    "from torch_geometric.nn.inits import glorot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(42)\n",
    "random.seed(42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 3\n",
    "N_h = 2\n",
    "D_in  = 4\n",
    "D_out = 8\n",
    "\n",
    "edge_index = torch.tensor([[0, 0], [0, 1], [0, 2], [1, 0], [1, 1], [2, 0], [2, 2]]).t()\n",
    "edge_src = edge_index[:1].t().squeeze()\n",
    "edge_tgt = edge_index[1:].t().squeeze()\n",
    "\n",
    "X = torch.randn((N, D_in))\n",
    "W = torch.randn((N_h, D_out, D_in))\n",
    "W_res = torch.randn((N_h, D_out, D_in))\n",
    "A_src = torch.randn((N_h, D_out, 1))\n",
    "A_tgt = torch.randn((N_h, D_out, 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "act = F.elu"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Manual forward pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([8, 4])\n",
      "torch.Size([8])\n",
      "torch.Size([8])\n"
     ]
    }
   ],
   "source": [
    "W_0 = W[0,:,:]\n",
    "a_src_0 = A_src[0,:].squeeze()\n",
    "a_tgt_0 = A_tgt[0,:].squeeze()\n",
    "print(W_0.shape)\n",
    "print(a_src_0.shape)\n",
    "print(a_tgt_0.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "H_in = X"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.0287, -0.9554, -0.1101, -0.3255,  0.2273,  0.1931,  0.1718, -0.2803],\n",
       "        [-1.5151, -3.7476,  2.5922, -1.3642,  3.9152, -3.7800,  0.0893, -3.4836],\n",
       "        [-0.0877, -1.9032, -0.3221, -0.4655,  0.2594, -0.0806,  1.4462, -1.3545]])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "H_w = H_in @ W_0.T\n",
    "H_w"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.0287, -0.9554, -0.1101, -0.3255,  0.2273,  0.1931,  0.1718, -0.2803],\n",
       "        [-1.5151, -3.7476,  2.5922, -1.3642,  3.9152, -3.7800,  0.0893, -3.4836],\n",
       "        [-0.0877, -1.9032, -0.3221, -0.4655,  0.2594, -0.0806,  1.4462, -1.3545]])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "H_w = torch.einsum('ij, kj -> ik', H_in, W_0)\n",
    "H_w # (|V|, D_out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True]])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.isclose(H_in @ W_0.T, torch.einsum('ij, kj -> ik', H_in, W_0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0, 0, 0, 1, 1, 2, 2]) tensor([0, 1, 2, 0, 1, 0, 2])\n"
     ]
    }
   ],
   "source": [
    "print(edge_src, edge_tgt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.0287, -0.9554, -0.1101, -0.3255,  0.2273,  0.1931,  0.1718, -0.2803],\n",
       "        [-0.0287, -0.9554, -0.1101, -0.3255,  0.2273,  0.1931,  0.1718, -0.2803],\n",
       "        [-0.0287, -0.9554, -0.1101, -0.3255,  0.2273,  0.1931,  0.1718, -0.2803],\n",
       "        [-1.5151, -3.7476,  2.5922, -1.3642,  3.9152, -3.7800,  0.0893, -3.4836],\n",
       "        [-1.5151, -3.7476,  2.5922, -1.3642,  3.9152, -3.7800,  0.0893, -3.4836],\n",
       "        [-0.0877, -1.9032, -0.3221, -0.4655,  0.2594, -0.0806,  1.4462, -1.3545],\n",
       "        [-0.0877, -1.9032, -0.3221, -0.4655,  0.2594, -0.0806,  1.4462, -1.3545]])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "H_w_src = torch.index_select(H_w, 0, edge_src) # (|E|, D_out)\n",
    "H_w_src"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.0287, -0.9554, -0.1101, -0.3255,  0.2273,  0.1931,  0.1718, -0.2803],\n",
       "        [-1.5151, -3.7476,  2.5922, -1.3642,  3.9152, -3.7800,  0.0893, -3.4836],\n",
       "        [-0.0877, -1.9032, -0.3221, -0.4655,  0.2594, -0.0806,  1.4462, -1.3545],\n",
       "        [-0.0287, -0.9554, -0.1101, -0.3255,  0.2273,  0.1931,  0.1718, -0.2803],\n",
       "        [-1.5151, -3.7476,  2.5922, -1.3642,  3.9152, -3.7800,  0.0893, -3.4836],\n",
       "        [-0.0287, -0.9554, -0.1101, -0.3255,  0.2273,  0.1931,  0.1718, -0.2803],\n",
       "        [-0.0877, -1.9032, -0.3221, -0.4655,  0.2594, -0.0806,  1.4462, -1.3545]])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "H_w_tgt = torch.index_select(H_w, 0, edge_tgt) # (|E|, D_out)\n",
    "H_w_tgt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-0.1287, -0.1287, -0.1287, -5.6065, -5.6065,  0.0103,  0.0103])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "E_pre_src = H_w_src @ a_src_0 # (|E|, 1), a_src^T Whj\n",
    "E_pre_src"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-0.1287, -0.1287, -0.1287, -5.6065, -5.6065,  0.0103,  0.0103])"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "E_pre_src = torch.einsum('ij, j -> i', H_w_src, a_src_0)\n",
    "E_pre_src"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-0.4347, -4.3002, -0.8585, -0.4347, -4.3002, -0.4347, -0.8585])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "E_pre_tgt = H_w_tgt @ a_tgt_0 # (|E|, 1), a_tgt^T Whj\n",
    "E_pre_tgt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "E_pre = E_pre_src + E_pre_tgt  # (|E|, 1), a^T [Whi || Whj]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([7])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "E = F.leaky_relu(E_pre, negative_slope=0.2) # (|E|, 1), LeakyRelu(a^T [Whi || Whj])\n",
    "E.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.4233, 0.7494, 0.4931, 0.1415, 0.2506, 0.4352, 0.5069])"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "alpha_scores = softmax(E, edge_tgt, dim=0) # (|E|, 1)\n",
    "alpha_scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([1., 1., 1.])"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "scatter(alpha_scores, edge_tgt, dim=0, reduce='sum')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch_geometric.utils as U"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.4233, 0.1415, 0.4352],\n",
       "        [0.7494, 0.2506, 0.0000],\n",
       "        [0.4931, 0.0000, 0.5069]])"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Alpha = U.to_dense_adj(edge_index, edge_attr=alpha_scores)[0].t()  # alpha from paper\n",
    "Alpha # Alpha_{i, j} = attention paid to node j when computing embedding for node i"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.4233, 0.4233, 0.4233, 0.4233, 0.4233, 0.4233, 0.4233, 0.4233],\n",
       "        [0.7494, 0.7494, 0.7494, 0.7494, 0.7494, 0.7494, 0.7494, 0.7494],\n",
       "        [0.4931, 0.4931, 0.4931, 0.4931, 0.4931, 0.4931, 0.4931, 0.4931],\n",
       "        [0.1415, 0.1415, 0.1415, 0.1415, 0.1415, 0.1415, 0.1415, 0.1415],\n",
       "        [0.2506, 0.2506, 0.2506, 0.2506, 0.2506, 0.2506, 0.2506, 0.2506],\n",
       "        [0.4352, 0.4352, 0.4352, 0.4352, 0.4352, 0.4352, 0.4352, 0.4352],\n",
       "        [0.5069, 0.5069, 0.5069, 0.5069, 0.5069, 0.5069, 0.5069, 0.5069]])"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Alpha = alpha_scores.repeat(D_out, 1).t() # (|E|, D_out)\n",
    "Alpha"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.2326, -0.8285,  0.1801, -0.4134,  0.7632, -0.3863,  0.7148, -0.6991],\n",
       "        [-0.3304, -0.8089,  0.5670, -0.4433,  1.1514, -0.5518,  0.1511, -0.6614],\n",
       "        [-0.0569, -0.7621, -0.1955, -0.3273,  0.2436,  0.0544,  0.8179, -0.5617]])"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "H_out_pre = scatter(Alpha * H_w_src, edge_tgt, dim=0, reduce='sum')  # (|V|, D_out)\n",
    "H_out = act(H_out_pre) # (|V|, D_out)\n",
    "H_out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### GATHead implementation & tests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GATHead(nn.Module):\n",
    "  def __init__(self, D_in: int, D_out: int, act=F.elu, dropout:float =0.0):\n",
    "    super().__init__()\n",
    "    self.D_in = D_in\n",
    "    self.D_out = D_out\n",
    "    self.act = act\n",
    "    self.dropout = dropout\n",
    "\n",
    "    self.W = nn.Parameter(torch.zeros((D_out, D_in)))\n",
    "    self.a_src = nn.Parameter(torch.zeros((D_out, 1)))\n",
    "    self.a_tgt = nn.Parameter(torch.zeros((D_out, 1)))\n",
    "\n",
    "    self.reset_parameters()\n",
    "\n",
    "  def reset_parameters(self):\n",
    "    glorot(self.W)\n",
    "    glorot(self.a_src)\n",
    "    glorot(self.a_tgt)\n",
    "\n",
    "  def forward(self, H_in: torch.tensor, edge_index: torch.tensor) -> torch.tensor:\n",
    "    edge_src = edge_index[:1].t().squeeze()\n",
    "    edge_tgt = edge_index[1:].t().squeeze()\n",
    "\n",
    "    D_out = self.D_out\n",
    "    W = self.W\n",
    "    a_src = self.a_src.squeeze()\n",
    "    a_tgt = self.a_tgt.squeeze()\n",
    "    act = self.act\n",
    "    dropout = self.dropout\n",
    "    training = self.training\n",
    "\n",
    "    H_in = F.dropout(H_in, dropout, training)\n",
    "\n",
    "    H_w = torch.einsum('ij, kj -> ik', H_in, W)  # (|V|, D_out)\n",
    "    H_w = F.dropout(H_w, dropout, training) # (|V|, D_out)\n",
    "\n",
    "    H_w_src = torch.index_select(H_w, 0, edge_src) # (|E|, D_out)\n",
    "    H_w_tgt = torch.index_select(H_w, 0, edge_tgt) # (|E|, D_out)\n",
    "    \n",
    "    E_pre_src = torch.einsum('ij, j -> i', H_w_src, a_src) # (|E|, 1), a_src^T Whi\n",
    "    E_pre_tgt = torch.einsum('ij, j -> i', H_w_tgt, a_tgt) # (|E|, 1), a_tgt^T Whj\n",
    "\n",
    "    E_pre = E_pre_src + E_pre_tgt  # (|E|, 1), a^T [Whi || Whj]\n",
    "    E = F.leaky_relu(E_pre, negative_slope=0.2) # (|E|, 1), LeakyRelu(a^T [Whi || Whj])\n",
    "\n",
    "    alpha_scores = softmax(E, edge_tgt) # (|E|, 1)\n",
    "    alpha_scores = F.dropout(alpha_scores, dropout, training)\n",
    "\n",
    "    Alpha = alpha_scores.repeat(D_out, 1).t() # (|E|, D_out)\n",
    "\n",
    "    H_out_pre = scatter(Alpha * H_w_src, edge_tgt, dim=0, reduce='sum') # (|V|, D_out)\n",
    "    H_out = act(H_out_pre) # (|V|, D_out)\n",
    "    return H_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "gat_head = GATHead(D_in, D_out, F.elu)\n",
    "\n",
    "W_0 = W[0,:,:]\n",
    "a_src_0 = A_src[0,:]\n",
    "a_tgt_0 = A_tgt[0,:]\n",
    "\n",
    "gat_head.a_src = nn.Parameter(a_src_0, requires_grad=False)\n",
    "gat_head.a_tgt = nn.Parameter(a_tgt_0, requires_grad=False)\n",
    "gat_head.W = nn.Parameter(W_0, requires_grad=False)\n",
    "\n",
    "assert(torch.equal(gat_head(X, edge_index), H_out))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "gat_head = GATHead(D_in, D_out, nn.Identity())\n",
    "\n",
    "W_0 = W[0,:,:]\n",
    "a_src_0 = A_src[0,:]\n",
    "a_tgt_0 = A_tgt[0,:]\n",
    "\n",
    "gat_head.a_src = nn.Parameter(a_src_0, requires_grad=False)\n",
    "gat_head.a_tgt = nn.Parameter(a_tgt_0, requires_grad=False)\n",
    "gat_head.W = nn.Parameter(W_0, requires_grad=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "head_gconv = gnn.GATConv(in_channels=D_in, out_channels=D_out, num_heads=1, bias=False, add_self_loops=False)\n",
    "\n",
    "head_gconv.att_src = nn.Parameter(a_src_0.squeeze(), requires_grad=False)\n",
    "head_gconv.att_dst = nn.Parameter(a_tgt_0.squeeze(), requires_grad=False)\n",
    "head_gconv.lin_src.weight = nn.Parameter(W_0, requires_grad=False)\n",
    "head_gconv.lin_dst = head_gconv.lin_src"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.2647, -1.7630,  0.1801, -0.5335,  0.7632, -0.4883,  0.7148, -1.2011],\n",
       "        [-0.4011, -1.6550,  0.5670, -0.5858,  1.1514, -0.8024,  0.1511, -1.0830],\n",
       "        [-0.0586, -1.4359, -0.2176, -0.3965,  0.2436,  0.0544,  0.8179, -0.8249]])"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gat_head(X, edge_index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.2647, -1.7630,  0.1801, -0.5335,  0.7632, -0.4883,  0.7148, -1.2011],\n",
       "        [-0.4011, -1.6550,  0.5670, -0.5858,  1.1514, -0.8024,  0.1511, -1.0830],\n",
       "        [-0.0586, -1.4359, -0.2176, -0.3965,  0.2436,  0.0544,  0.8179, -0.8249]])"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "head_gconv(X, edge_index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True]])"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.isclose(gat_head(X, edge_index), head_gconv(X, edge_index))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "head_gconv = gnn.GATConv(in_channels=D_in, out_channels=D_out, num_heads=1, bias=False, add_self_loops=False)\n",
    "\n",
    "head_gconv.att_src = nn.Parameter(a_src_0.squeeze(), requires_grad=False)\n",
    "head_gconv.att_dst = nn.Parameter(a_tgt_0.squeeze(), requires_grad=False)\n",
    "head_gconv.lin_src.weight = nn.Parameter(W_0, requires_grad=False)\n",
    "head_gconv.lin_dst = head_gconv.lin_src"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[-0.2647, -1.7630,  0.1801, -0.5335,  0.7632, -0.4883,  0.7148, -1.2011],\n",
       "         [-0.4011, -1.6550,  0.5670, -0.5858,  1.1514, -0.8024,  0.1511, -1.0830],\n",
       "         [-0.0586, -1.4359, -0.2176, -0.3965,  0.2436,  0.0544,  0.8179, -0.8249]]),\n",
       " (tensor([[0, 0, 0, 1, 1, 2, 2],\n",
       "          [0, 1, 2, 0, 1, 0, 2]]),\n",
       "  tensor([[0.4233],\n",
       "          [0.7494],\n",
       "          [0.4931],\n",
       "          [0.1415],\n",
       "          [0.2506],\n",
       "          [0.4352],\n",
       "          [0.5069]])))"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "head_gconv(X, edge_index, return_attention_weights=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### GATLayer implementation & tests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "from layers import GATLayer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Test whether GATLayer output matches the output of separate GATHeads"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "layer = GATLayer(D_in, D_out, 2, act=nn.Identity())\n",
    "layer.W = nn.Parameter(W, requires_grad=False)\n",
    "layer.A_src = nn.Parameter(A_src, requires_grad=False)\n",
    "layer.A_tgt = nn.Parameter(A_tgt, requires_grad=False)\n",
    "layer(X, edge_index)\n",
    "\n",
    "gat_head_0 = GATHead(D_in, D_out, nn.Identity())\n",
    "\n",
    "W_0 = W[0,:,:]\n",
    "a_src_0 = A_src[0,:]\n",
    "a_tgt_0 = A_tgt[0,:]\n",
    "\n",
    "gat_head_0.a_src = nn.Parameter(a_src_0, requires_grad=False)\n",
    "gat_head_0.a_tgt = nn.Parameter(a_tgt_0, requires_grad=False)\n",
    "gat_head_0.W = nn.Parameter(W_0, requires_grad=False)\n",
    "\n",
    "gat_head_1 = GATHead(D_in, D_out, nn.Identity())\n",
    "\n",
    "W_1 = W[1,:,:]\n",
    "a_src_1 = A_src[1,:]\n",
    "a_tgt_1 = A_tgt[1,:]\n",
    "\n",
    "gat_head_1.a_src = nn.Parameter(a_src_1, requires_grad=False)\n",
    "gat_head_1.a_tgt = nn.Parameter(a_tgt_1, requires_grad=False)\n",
    "gat_head_1.W = nn.Parameter(W_1, requires_grad=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[-0.2647, -1.7630,  0.1801, -0.5335,  0.7632, -0.4883,  0.7148,\n",
      "          -1.2011],\n",
      "         [-0.4011, -1.6550,  0.5670, -0.5858,  1.1514, -0.8024,  0.1511,\n",
      "          -1.0830],\n",
      "         [-0.0586, -1.4359, -0.2176, -0.3965,  0.2436,  0.0544,  0.8179,\n",
      "          -0.8249]],\n",
      "\n",
      "        [[ 0.4504,  0.5838, -0.5094, -0.1251,  0.0229,  1.1117,  1.4967,\n",
      "           0.2324],\n",
      "         [ 0.0923,  0.3843, -0.4492,  0.1707,  0.0470,  0.3850,  1.0480,\n",
      "          -0.0914],\n",
      "         [ 0.3780,  0.5470, -0.4937, -0.0965,  0.0248,  1.1349,  1.5795,\n",
      "           0.2952]]])\n",
      "tensor([[[-0.2647, -1.7630,  0.1801, -0.5335,  0.7632, -0.4883,  0.7148,\n",
      "          -1.2011],\n",
      "         [-0.4011, -1.6550,  0.5670, -0.5858,  1.1514, -0.8024,  0.1511,\n",
      "          -1.0830],\n",
      "         [-0.0586, -1.4359, -0.2176, -0.3965,  0.2436,  0.0544,  0.8179,\n",
      "          -0.8249]],\n",
      "\n",
      "        [[ 0.4504,  0.5838, -0.5094, -0.1251,  0.0229,  1.1117,  1.4967,\n",
      "           0.2324],\n",
      "         [ 0.0923,  0.3843, -0.4492,  0.1707,  0.0470,  0.3850,  1.0480,\n",
      "          -0.0914],\n",
      "         [ 0.3780,  0.5470, -0.4937, -0.0965,  0.0248,  1.1349,  1.5795,\n",
      "           0.2952]]])\n"
     ]
    }
   ],
   "source": [
    "H_out_head_0 = gat_head_0(X, edge_index)\n",
    "H_out_head_1 = gat_head_1(X, edge_index)\n",
    "\n",
    "H_out_heads = torch.cat((H_out_head_0, H_out_head_1), 0).view((2, N, D_out))\n",
    "H_out_layer = layer(X, edge_index)\n",
    "\n",
    "print(H_out_heads)\n",
    "print(H_out_layer)\n",
    "\n",
    "assert(torch.equal(H_out_heads, H_out_layer))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test whether GATLayer output matches the output of pytorch geometric GATConv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Concat mode, identity activation test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "layer = GATLayer(D_in, D_out, 2, act=nn.Identity(), reduce='concat')\n",
    "layer.W = nn.Parameter(W, requires_grad=False)\n",
    "layer.A_src = nn.Parameter(A_src, requires_grad=False)\n",
    "layer.A_tgt = nn.Parameter(A_tgt, requires_grad=False)\n",
    "\n",
    "layer_gconv = gnn.GATConv(in_channels=D_in, out_channels=D_out, heads=2, bias=False, add_self_loops=False)\n",
    "layer_gconv.att_src = nn.Parameter(A_src.view((1, N_h, D_out)), requires_grad=False)\n",
    "layer_gconv.att_dst = nn.Parameter(A_tgt.view((1, N_h, D_out)), requires_grad=False)\n",
    "layer_gconv.lin_src.weight = nn.Parameter(W.view((N_h * D_out, D_in)), requires_grad=False)\n",
    "layer_gconv.lin_dst = head_gconv.lin_src"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[-0.2647, -1.7630,  0.1801, -0.5335,  0.7632, -0.4883,  0.7148, -1.2011,\n",
      "          0.4504,  0.5838, -0.5094, -0.1251,  0.0229,  1.1117,  1.4967,  0.2324],\n",
      "        [-0.4011, -1.6550,  0.5670, -0.5858,  1.1514, -0.8024,  0.1511, -1.0830,\n",
      "          0.0923,  0.3843, -0.4492,  0.1707,  0.0470,  0.3850,  1.0480, -0.0914],\n",
      "        [-0.0586, -1.4359, -0.2176, -0.3965,  0.2436,  0.0544,  0.8179, -0.8249,\n",
      "          0.3780,  0.5470, -0.4937, -0.0965,  0.0248,  1.1349,  1.5795,  0.2952]])\n",
      "tensor([[-0.2647, -1.7630,  0.1801, -0.5335,  0.7632, -0.4883,  0.7148, -1.2011,\n",
      "          0.4504,  0.5838, -0.5094, -0.1251,  0.0229,  1.1117,  1.4967,  0.2324],\n",
      "        [-0.4011, -1.6550,  0.5670, -0.5858,  1.1514, -0.8024,  0.1511, -1.0830,\n",
      "          0.0923,  0.3843, -0.4492,  0.1707,  0.0470,  0.3850,  1.0480, -0.0914],\n",
      "        [-0.0586, -1.4359, -0.2176, -0.3965,  0.2436,  0.0544,  0.8179, -0.8249,\n",
      "          0.3780,  0.5470, -0.4937, -0.0965,  0.0248,  1.1349,  1.5795,  0.2952]])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True]])"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out_layer = layer(X, edge_index)\n",
    "out_layer_gconv = layer_gconv(X, edge_index)\n",
    "\n",
    "print(out_layer)\n",
    "print(out_layer_gconv)\n",
    "\n",
    "torch.isclose(out_layer, out_layer_gconv)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Concat mode, ELU activation test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "layer = GATLayer(D_in, D_out, 2, act=F.elu, reduce='concat')\n",
    "layer.W = nn.Parameter(W, requires_grad=False)\n",
    "layer.A_src = nn.Parameter(A_src, requires_grad=False)\n",
    "layer.A_tgt = nn.Parameter(A_tgt, requires_grad=False)\n",
    "\n",
    "layer_gconv = gnn.GATConv(in_channels=D_in, out_channels=D_out, heads=2, bias=False, add_self_loops=False)\n",
    "layer_gconv.att_src = nn.Parameter(A_src.view((1, N_h, D_out)), requires_grad=False)\n",
    "layer_gconv.att_dst = nn.Parameter(A_tgt.view((1, N_h, D_out)), requires_grad=False)\n",
    "layer_gconv.lin_src.weight = nn.Parameter(W.view((N_h * D_out, D_in)), requires_grad=False)\n",
    "layer_gconv.lin_dst = head_gconv.lin_src"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[-0.2326, -0.8285,  0.1801, -0.4134,  0.7632, -0.3863,  0.7148, -0.6991,\n",
      "          0.4504,  0.5838, -0.3991, -0.1176,  0.0229,  1.1117,  1.4967,  0.2324],\n",
      "        [-0.3304, -0.8089,  0.5670, -0.4433,  1.1514, -0.5518,  0.1511, -0.6614,\n",
      "          0.0923,  0.3843, -0.3619,  0.1707,  0.0470,  0.3850,  1.0480, -0.0874],\n",
      "        [-0.0569, -0.7621, -0.1955, -0.3273,  0.2436,  0.0544,  0.8179, -0.5617,\n",
      "          0.3780,  0.5470, -0.3896, -0.0919,  0.0248,  1.1349,  1.5795,  0.2952]])\n",
      "tensor([[-0.2326, -0.8285,  0.1801, -0.4134,  0.7632, -0.3863,  0.7148, -0.6991,\n",
      "          0.4504,  0.5838, -0.3991, -0.1176,  0.0229,  1.1117,  1.4967,  0.2324],\n",
      "        [-0.3304, -0.8089,  0.5670, -0.4433,  1.1514, -0.5518,  0.1511, -0.6614,\n",
      "          0.0923,  0.3843, -0.3619,  0.1707,  0.0470,  0.3850,  1.0480, -0.0874],\n",
      "        [-0.0569, -0.7621, -0.1955, -0.3273,  0.2436,  0.0544,  0.8179, -0.5617,\n",
      "          0.3780,  0.5470, -0.3896, -0.0919,  0.0248,  1.1349,  1.5795,  0.2952]])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True]])"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out_layer = layer(X, edge_index)\n",
    "out_layer_gconv = F.elu(layer_gconv(X, edge_index))\n",
    "\n",
    "print(out_layer)\n",
    "print(out_layer_gconv)\n",
    "\n",
    "torch.isclose(out_layer, out_layer_gconv)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Avg mode, identity activation test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "layer = GATLayer(D_in, D_out, 2, act=nn.Identity(), reduce='avg')\n",
    "layer.W = nn.Parameter(W, requires_grad=False)\n",
    "layer.A_src = nn.Parameter(A_src, requires_grad=False)\n",
    "layer.A_tgt = nn.Parameter(A_tgt, requires_grad=False)\n",
    "\n",
    "layer_gconv = gnn.GATConv(in_channels=D_in, out_channels=D_out, heads=2, bias=False, add_self_loops=False, concat=False)\n",
    "layer_gconv.att_src = nn.Parameter(A_src.view((1, N_h, D_out)), requires_grad=False)\n",
    "layer_gconv.att_dst = nn.Parameter(A_tgt.view((1, N_h, D_out)), requires_grad=False)\n",
    "layer_gconv.lin_src.weight = nn.Parameter(W.view((N_h * D_out, D_in)), requires_grad=False)\n",
    "layer_gconv.lin_dst = head_gconv.lin_src"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[ 0.0928, -0.5896, -0.1647, -0.3293,  0.3930,  0.3117,  1.1057, -0.4844],\n",
      "        [-0.1544, -0.6354,  0.0589, -0.2075,  0.5992, -0.2087,  0.5996, -0.5872],\n",
      "        [ 0.1597, -0.4444, -0.3556, -0.2465,  0.1342,  0.5946,  1.1987, -0.2648]])\n",
      "tensor([[ 0.0928, -0.5896, -0.1647, -0.3293,  0.3930,  0.3117,  1.1057, -0.4844],\n",
      "        [-0.1544, -0.6354,  0.0589, -0.2075,  0.5992, -0.2087,  0.5996, -0.5872],\n",
      "        [ 0.1597, -0.4444, -0.3556, -0.2465,  0.1342,  0.5946,  1.1987, -0.2648]])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True]])"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out_layer = layer(X, edge_index)\n",
    "out_layer_gconv = layer_gconv(X, edge_index)\n",
    "\n",
    "print(out_layer)\n",
    "print(out_layer_gconv)\n",
    "\n",
    "torch.isclose(out_layer, out_layer_gconv)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Avg mode, softmax activation test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "layer = GATLayer(D_in, D_out, 2, act=nn.Softmax(dim=1), reduce='avg')\n",
    "layer.W = nn.Parameter(W, requires_grad=False)\n",
    "layer.A_src = nn.Parameter(A_src, requires_grad=False)\n",
    "layer.A_tgt = nn.Parameter(A_tgt, requires_grad=False)\n",
    "\n",
    "layer_gconv = gnn.GATConv(in_channels=D_in, out_channels=D_out, heads=2, bias=False, add_self_loops=False, concat=False)\n",
    "layer_gconv.att_src = nn.Parameter(A_src.view((1, N_h, D_out)), requires_grad=False)\n",
    "layer_gconv.att_dst = nn.Parameter(A_tgt.view((1, N_h, D_out)), requires_grad=False)\n",
    "layer_gconv.lin_src.weight = nn.Parameter(W.view((N_h * D_out, D_in)), requires_grad=False)\n",
    "layer_gconv.lin_dst = head_gconv.lin_src"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.1131, 0.0571, 0.0874, 0.0741, 0.1527, 0.1407, 0.3114, 0.0635],\n",
      "        [0.1036, 0.0641, 0.1283, 0.0983, 0.2202, 0.0981, 0.2202, 0.0672],\n",
      "        [0.1135, 0.0620, 0.0678, 0.0756, 0.1106, 0.1753, 0.3208, 0.0742]])\n",
      "tensor([[0.1131, 0.0571, 0.0874, 0.0741, 0.1527, 0.1407, 0.3114, 0.0635],\n",
      "        [0.1036, 0.0641, 0.1283, 0.0983, 0.2202, 0.0981, 0.2202, 0.0672],\n",
      "        [0.1135, 0.0620, 0.0678, 0.0756, 0.1106, 0.1753, 0.3208, 0.0742]])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True]])"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out_layer = layer(X, edge_index)\n",
    "out_layer_gconv = nn.Softmax(dim=1)(layer_gconv(X, edge_index))\n",
    "\n",
    "print(out_layer)\n",
    "print(out_layer_gconv)\n",
    "\n",
    "torch.isclose(out_layer, out_layer_gconv)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.0 ('grl-env': venv)",
   "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.0"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "4737062b79cbd82e7f605897a0c10fa5613d72f803a33da3548714c1d9e40d67"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}