fix doc

pyg-team · Jun 29, 2019 · 3bbfe02 · 3bbfe02
1 parent 57a8993
commit 3bbfe02
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/torch_geometric/nn/conv/gmm_conv.py b/torch_geometric/nn/conv/gmm_conv.py
@@ -13,8 +13,8 @@ class GMMConv(MessagePassing):
     <https://arxiv.org/abs/1611.08402>`_ paper
 
     .. math::
-        \mathbf{x}^{\prime}_i =
-        \sum_{j \in \mathcal{N}(i) \cup \{ i \}} \frac{1}{K} \sum_{k=1}^K
+        \mathbf{x}^{\prime}_i = \frac{1}{|\mathcal{N}(i)|}
+        \sum_{j \in \mathcal{N}(i) \frac{1}{K} \sum_{k=1}^K
         \mathbf{w}_k(\mathbf{e}_{i,j}) \odot \mathbf{\Theta}_k \mathbf{x}_j,
 
     where

diff --git a/torch_geometric/nn/pool/sag_pool.py b/torch_geometric/nn/pool/sag_pool.py
@@ -9,9 +9,12 @@ class SAGPooling(torch.nn.Module):
 
     .. math::
         \mathbf{y} &= \textrm{GNN}(\mathbf{X}, \mathbf{A})
+
         \mathbf{i} &= \mathrm{top}_k(\mathbf{y})
+
         \mathbf{X}^{\prime} &= (\mathbf{X} \odot
         \mathrm{tanh}(\mathbf{y}))_{\mathbf{i}}
+
         \mathbf{A}^{\prime} &= \mathbf{A}_{\mathbf{i},\mathbf{i}},
 
     where nodes are dropped based on a learnable projection score