[NN] Refactor the Code Structure of GT (dmlc#5100)

DominikaJedynak · Mar 12, 2024 · 696b84d · 696b84d
1 parent d753aa3
commit 696b84d
Show file tree

Hide file tree

Showing 17 changed files with 975 additions and 931 deletions.
diff --git a/docs/source/api/python/dgl.rst b/docs/source/api/python/dgl.rst
@@ -111,7 +111,7 @@ Operators for generating positional encodings of each node.
     :toctree: ../../generated
 
     random_walk_pe
-    laplacian_pe
+    lap_pe
     double_radius_node_labeling
     shortest_dist
     svd_pe

diff --git a/docs/source/api/python/nn-pytorch.rst b/docs/source/api/python/nn-pytorch.rst
@@ -146,3 +146,18 @@ Network Embedding Modules
 
     ~dgl.nn.pytorch.DeepWalk
     ~dgl.nn.pytorch.MetaPath2Vec
+
+Utility Modules for Graph Transformer
+----------------------------------------
+.. autosummary::
+    :toctree: ../../generated/
+    :nosignatures:
+    :template: classtemplate.rst
+
+    ~dgl.nn.pytorch.gt.DegreeEncoder
+    ~dgl.nn.pytorch.gt.LapPosEncoder
+    ~dgl.nn.pytorch.gt.PathEncoder
+    ~dgl.nn.pytorch.gt.SpatialEncoder
+    ~dgl.nn.pytorch.gt.SpatialEncoder3d
+    ~dgl.nn.pytorch.gt.BiasedMHA
+    ~dgl.nn.pytorch.gt.GraphormerLayer
diff --git a/docs/source/api/python/transforms.rst b/docs/source/api/python/transforms.rst
@@ -29,7 +29,7 @@ dgl.transforms
     DropEdge
     AddEdge
     RandomWalkPE
-    LaplacianPE
+    LapPE
     FeatMask
     RowFeatNormalizer
     SIGNDiffusion

diff --git a/python/dgl/nn/pytorch/__init__.py b/python/dgl/nn/pytorch/__init__.py
@@ -8,12 +8,6 @@
 from .factory import *
 from .hetero import *
 from .sparse_emb import NodeEmbedding
-from .utils import (
-    JumpingKnowledge,
-    LabelPropagation,
-    LaplacianPosEnc,
-    Sequential,
-    WeightBasis,
-)
+from .utils import JumpingKnowledge, LabelPropagation, Sequential, WeightBasis
 from .network_emb import *
-from .graph_transformer import *
+from .gt import *