PR [#76] nas -> dev

Add GASSO and AutoAttend

LICENSE

@@ -187,7 +187,7 @@
       same "printed page" as the copyright notice for easier
       identification within third-party archives.
 
-   Copyright [yyyy] [name of copyright owner]
+   Copyright 2020-2021 AGLTeam THUMNLab and contributors
 
    Licensed under the Apache License, Version 2.0 (the "License");
    you may not use this file except in compliance with the License.
@@ -199,4 +199,4 @@
    distributed under the License is distributed on an "AS IS" BASIS,
    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    See the License for the specific language governing permissions and
-   limitations under the License.
+   limitations under the License.

autogl/module/nas/algorithm/__init__.py

@@ -29,6 +29,9 @@ from .darts import Darts
 from .enas import Enas
 from .random_search import RandomSearch
 from .rl import RL, GraphNasRL
+from ..backend import *
+if not is_dgl():
+    from .gasso import Gasso
 from .spos import Spos
 
 def build_nas_algo_from_name(name: str) -> BaseNAS:
@@ -53,3 +56,5 @@ def build_nas_algo_from_name(name: str) -> BaseNAS:
 
 
 __all__ = ["BaseNAS", "Darts", "Enas", "RandomSearch", "RL", "GraphNasRL","Spos"]
+if not is_dgl():
+    __all__.append("Gasso")

autogl/module/nas/algorithm/darts.py

@@ -102,7 +102,7 @@ class Darts(BaseNAS):
         model_wd=5e-4,
         arch_lr=3e-4,
         arch_wd=1e-3,
-        device="cuda",
+        device="auto",
     ):
         super().__init__(device=device)
         self.num_epochs = num_epochs

autogl/module/nas/algorithm/enas.py

@@ -79,7 +79,7 @@ class Enas(BaseNAS):
         model_lr=5e-3,
         model_wd=5e-4,
         disable_progress=True,
-        device="cuda",
+        device="auto",
     ):
         super().__init__(device)
         self.device = device

autogl/module/nas/algorithm/gasso.py

@@ -0,0 +1,158 @@
+# "Graph differentiable architecture search with structure optimization" NeurIPS 21'
+
+import logging
+
+import torch
+import torch.optim
+import torch.nn as nn
+import torch.nn.functional as F
+
+from . import register_nas_algo
+from .base import BaseNAS
+from ..estimator.base import BaseEstimator
+from ..space import BaseSpace
+from ..utils import replace_layer_choice, replace_input_choice
+from ...model.base import BaseAutoModel
+
+from torch.autograd import Variable
+import numpy as np
+import time
+import copy
+import torch.optim as optim
+import scipy.sparse as sp
+
+_logger = logging.getLogger(__name__)
+
+@register_nas_algo("gasso")
+class Gasso(BaseNAS):
+    """
+    GASSO trainer.
+
+    Parameters
+    ----------
+    num_epochs : int
+        Number of epochs planned for training.
+    warmup_epochs : int
+        Number of epochs planned for warming up.
+    workers : int
+        Workers for data loading.
+    model_lr : float
+        Learning rate to optimize the model.
+    model_wd : float
+        Weight decay to optimize the model.
+    arch_lr : float
+        Learning rate to optimize the architecture.
+    stru_lr : float
+        Learning rate to optimize the structure.
+    lamb : float
+        The parameter to control the influence of hidden feature smoothness
+    device : str or torch.device
+        The device of the whole process
+    """
+    def __init__(
+        self,
+        num_epochs=250,
+        warmup_epochs=10,
+        model_lr=0.01,
+        model_wd=1e-4,
+        arch_lr = 0.03,
+        stru_lr = 0.04,
+        lamb = 0.6,
+        device="auto",
+    ):
+        super().__init__(device=device)
+        self.device = device
+        self.num_epochs = num_epochs
+        self.warmup_epochs = warmup_epochs
+        self.model_lr = model_lr
+        self.model_wd = model_wd
+        self.arch_lr = arch_lr
+        self.stru_lr = stru_lr
+        self.lamb = lamb
+
+    def train_stru(self, model, optimizer, data):
+        # forward
+        model.train()
+        data[0].adj = self.adjs
+        logits = model(data[0]).detach()
+        loss = 0
+        for adj in self.adjs:
+            e1 = adj[0][0]
+            e2 = adj[0][1]
+            ew = adj[1]
+            diff = (logits[e1] - logits[e2]).pow(2).sum(1)
+            smooth = (diff * torch.sigmoid(ew)).sum()
+            dist = (ew * ew).sum()
+            loss += self.lamb * smooth + dist
+            
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+        train_loss = loss.item()
+        del logits
+
+    def _infer(self, model: BaseSpace, dataset, estimator: BaseEstimator, mask="train"):
+        dataset[0].adj = self.adjs
+        metric, loss = estimator.infer(model, dataset, mask=mask)
+        return metric, loss
+
+    def prepare(self, dset):
+        """Train Pro-GNN.
+        """
+        data = dset[0]
+        self.ews = []
+        self.edges = data.edge_index.to(self.device)
+        edge_weight = torch.ones(self.edges.size(1)).to(self.device)
+
+        self.adjs = []
+        for i in range(self.steps):
+            edge_weight = Variable(edge_weight * 1.0, requires_grad = True).to(self.device)
+            self.ews.append(edge_weight)
+            self.adjs.append((self.edges, edge_weight))
+
+    def fit(self, data):
+        self.optimizer = optim.Adam(self.space.parameters(), lr=self.model_lr, weight_decay=self.model_wd)
+        self.arch_optimizer = optim.Adam(self.space.arch_parameters(),
+                                              lr=self.arch_lr, betas=(0.5, 0.999))
+        self.stru_optimizer = optim.SGD(self.ews, lr=self.stru_lr)
+
+        # Train model
+        best_performance = 0
+        min_val_loss = float("inf")
+        min_train_loss = float("inf")
+
+        t_total = time.time()
+        for epoch in range(self.num_epochs):
+            self.space.train()
+            self.optimizer.zero_grad()
+            _, loss = self._infer(self.space, data, self.estimator, "train")
+            loss.backward()
+            self.optimizer.step()
+
+            if epoch <20:
+                continue
+            self.train_stru(self.space, self.stru_optimizer, data)
+            
+            self.arch_optimizer.zero_grad()
+            _, loss = self._infer(self.space, data, self.estimator, "train")
+            loss.backward()
+            self.arch_optimizer.step()
+
+            self.space.eval()
+            train_acc, _ = self._infer(self.space, data, self.estimator, "train")
+            val_acc, val_loss = self._infer(self.space, data, self.estimator, "val")
+            if val_loss < min_val_loss:
+                min_val_loss = val_loss
+                best_performance = val_acc
+                self.space.keep_prediction()
+            #print("acc:" + str(train_acc) + " val_acc" + str(val_acc))
+
+        return best_performance, min_val_loss
+
+    def search(self, space: BaseSpace, dataset, estimator):
+        self.estimator = estimator
+        self.space = space.to(self.device)
+        self.steps = space.steps
+        self.prepare(dataset)
+        perf, val_loss = self.fit(dataset)
+        return space.parse_model(None, self.device)

autogl/module/nas/algorithm/random_search.py

@@ -35,7 +35,7 @@ class RandomSearch(BaseNAS):
         Control whether show the progress bar.
     """
 
-    def __init__(self, device="cuda", num_epochs=400, disable_progress=False, hardware_metric_limit=None):
+    def __init__(self, device="auto", num_epochs=400, disable_progress=False, hardware_metric_limit=None):
         super().__init__(device)
         self.num_epochs = num_epochs
         self.disable_progress = disable_progress

autogl/module/nas/algorithm/rl.py

@@ -250,7 +250,7 @@ class RL(BaseNAS):
     def __init__(
         self,
         num_epochs=5,
-        device="cuda",
+        device="auto",
         log_frequency=None,
         grad_clip=5.0,
         entropy_weight=0.0001,
@@ -429,7 +429,7 @@ class GraphNasRL(BaseNAS):
 
     def __init__(
         self,
-        device="cuda",
+        device="auto",
         num_epochs=10,
         log_frequency=None,
         grad_clip=5.0,
@@ -580,6 +580,8 @@ class GraphNasRL(BaseNAS):
                 self.ctrl_optim.step()
 
                 bar.set_postfix(acc=metric, max_acc=max(rewards))
+        
+        LOGGER.info(f"epoch:{}, mean rewards:{}".format(epoch, sum(rewards) / len(rewards)))
         return sum(rewards) / len(rewards)
 
     def _resample(self):

autogl/module/nas/estimator/one_shot.py

@@ -1,4 +1,5 @@
 import torch.nn.functional as F
+import torch
 
 from . import register_nas_estimator
 from ..space import BaseSpace
@@ -7,13 +8,19 @@ from ..backend import *
 from ...train.evaluation import Acc
 from ..utils import get_hardware_aware_metric
 
-
 @register_nas_estimator("oneshot")
 class OneShotEstimator(BaseEstimator):
     """
     One shot estimator.
 
     Use model directly to get estimations.
+
+    Parameters
+    ----------
+    loss_f : str
+        The name of loss funciton in PyTorch
+    evaluation : list of Evaluation
+        The evaluation metrics in module/train/evaluation
     """
 
     def __init__(self, loss_f="nll_loss", evaluation=[Acc()]):
@@ -31,6 +38,7 @@ class OneShotEstimator(BaseEstimator):
 
         loss = getattr(F, self.loss_f)(pred, y)
         probs = F.softmax(pred, dim=1).detach().cpu().numpy()
+        
         y = y.cpu()
         metrics = [eva.evaluate(probs, y) for eva in self.evaluation]
         return metrics, loss
@@ -41,7 +49,18 @@ class OneShotEstimator_HardwareAware(OneShotEstimator):
     """
     One shot hardware-aware estimator.
 
-    Use model directly to get estimations.
+    Use model directly to get estimations with some hardware-aware metrics.
+
+    Parameters
+    ----------
+    loss_f : str
+        The name of loss funciton in PyTorch
+    evaluation : list of Evaluation
+        The evaluation metrics in module/train/evaluation
+    hardware_evaluation : str or runable
+        The hardware-aware metrics. Can be 'parameter' or 'latency'. Or you can define a special metric by a runable function
+    hardware_metric_weight : float
+        The weight of hardware-aware metric, which will be a bias added to metrics
     """
 
     def __init__(

autogl/module/nas/estimator/train_scratch.py

@@ -11,6 +11,13 @@ from autogl.module.train import NodeClassificationFullTrainer, Acc
 class TrainEstimator(BaseEstimator):
     """
     An estimator which trans from scratch
+
+    Parameters
+    ----------
+    loss_f : str
+        The name of loss funciton in PyTorch
+    evaluation : list of Evaluation
+        The evaluation metrics in module/train/evaluation
     """
 
     def __init__(self, loss_f="nll_loss", evaluation=[Acc()]):

autogl/module/nas/space/__init__.py

@@ -23,6 +23,9 @@ from .graph_nas_macro import GraphNasMacroNodeClassificationSpace
 from .graph_nas import GraphNasNodeClassificationSpace
 from .single_path import SinglePathNodeClassificationSpace
 
+from ..backend import *
+if not is_dgl():
+    from .gasso import GassoSpace
 
 def build_nas_space_from_name(name: str) -> BaseSpace:
     """
@@ -51,3 +54,6 @@ __all__ = [
     "GraphNasNodeClassificationSpace",
     "SinglePathNodeClassificationSpace",
 ]
+
+if not is_dgl():
+    __all__.append("GassoSpace")

autogl/module/nas/space/autoattend.py

@@ -0,0 +1,203 @@
+# codes in this file are reproduced from AutoAttend with some changes.
+from nni.nas.pytorch.mutables import Mutable
+import typing as _typ
+import torch
+
+import torch.nn.functional as F
+from nni.nas.pytorch import mutables
+
+from . import register_nas_space
+from .base import BaseSpace
+from ...model import BaseModel
+from ..utils import count_parameters, measure_latency
+
+from torch import nn
+from .operation import act_map, gnn_map
+
+from ..backend import *
+
+from .autoattend_space.ops1 import OPS as OPS1
+from .autoattend_space.ops2 import OPS as OPS2
+from .autoattend_space.operations import agg_map
+OPS = [OPS1, OPS2]
+
+
+@register_nas_space("autoattend")
+class AutoAttendNodeClassificationSpace(BaseSpace):
+    """
+    AutoAttend Search Space , please refer to http://proceedings.mlr.press/v139/guan21a.html for details.
+    The current implementation is nc (no context weight sharing), 
+    we will in future add other types of partial weight sharing proposed in the paper.
+
+    Parameters
+    ----------
+    ops_type : int 
+        0 or 1 , choosing from two sets of ops with index ops_type
+    gnn_ops : list of str
+        op names for searching, which descripts the compostion of operation pool
+    act_op : str
+        determine used activation function 
+    agg_ops : list of str
+        agg op names for searching. Only ['add','attn'] are options, as mentioned in the paper.
+    """
+    def __init__(
+        self,
+        hidden_dim: _typ.Optional[int] = 64,
+        layer_number: _typ.Optional[int] = 2,
+        dropout: _typ.Optional[float] = 0.9,
+        input_dim: _typ.Optional[int] = None,
+        output_dim: _typ.Optional[int] = None,
+        ops_type=0,
+        gnn_ops: _typ.Sequence[_typ.Union[str, _typ.Any]
+                               ] = None,
+        act_op="tanh",
+        head=8,
+        agg_ops=['add', 'attn']
+    ):
+        super().__init__()
+        self.layer_number = layer_number
+        self.hidden_dim = hidden_dim
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+        self.gnn_ops = gnn_ops
+        self.dropout = dropout
+        self.act_op = act_op
+        self.ops_type = ops_type
+        self.head = head
+        self.agg_ops = agg_ops
+
+    def instantiate(
+        self,
+        hidden_dim: _typ.Optional[int] = None,
+        layer_number: _typ.Optional[int] = None,
+        dropout: _typ.Optional[float] = None,
+        input_dim: _typ.Optional[int] = None,
+        output_dim: _typ.Optional[int] = None,
+        ops_type=None,
+        gnn_ops: _typ.Sequence[_typ.Union[str, _typ.Any]] = None,
+        act_op=None,
+        head=None,
+        agg_ops=None,
+        # con_ops: _typ.Sequence[_typ.Union[str, _typ.Any]] = None,
+    ):
+        super().instantiate()
+        self.dropout = dropout or self.dropout
+        self.hidden_dim = hidden_dim or self.hidden_dim
+        self.layer_number = layer_number or self.layer_number
+        self.input_dim = input_dim or self.input_dim
+        self.output_dim = output_dim or self.output_dim
+        self.gnn_ops = gnn_ops or self.gnn_ops
+        self.act_op = act_op or self.act_op
+        self.act = act_map(self.act_op)
+        self.head = head or self.head
+        self.ops_type = ops_type or self.ops_type
+        self.agg_ops = agg_ops or self.agg_ops
+        PRIMITIVES = list(OPS[self.ops_type].keys())
+        self.gnn_map = lambda x, * \
+            args, **kwargs: OPS[self.ops_type][x](*args, **kwargs)
+        self.gnn_ops = self.gnn_ops or PRIMITIVES
+        self.agg_map = lambda x, * \
+            args, **kwargs: agg_map[x](*args, **kwargs)
+        self.preproc0 = nn.Linear(self.input_dim, self.hidden_dim)
+        
+        node_labels = []
+        for layer in range(1, self.layer_number+1):
+            # stem path
+            key = f"stem_{layer}"
+            self._set_layer_choice(layer, key)
+
+            # side path
+            key = f"side_{layer}"
+            for i in range(2):
+                sub_key = f"{key}_{i}"
+                self._set_layer_choice(layer, sub_key)
+            node_labels.append(key)
+
+            # input
+            key = f"in_{layer}"
+            # self._set_input_choice(key,layer, choose_from=node_labels, n_chosen=1, return_mask=False)
+            self._set_input_choice(key, layer, n_candidates=len(
+                node_labels), n_chosen=1, return_mask=False)
+
+            # agg
+            key = f"agg_{layer}"
+            self._set_agg_choice(layer, key=key)
+
+        self._initialized = True
+
+        self.classifier2 = nn.Linear(self.hidden_dim, self.output_dim)
+
+    def _set_agg_choice(self, layer, key):
+        ops = [self.agg_map(op, self.hidden_dim, self.head,
+                            self.dropout)for op in self.agg_ops]
+        choice = self.setLayerChoice(
+            layer,
+            ops,
+            key=key,
+        )
+        setattr(self, key, choice)
+        return choice
+
+    def _set_layer_choice(self, layer, key):
+        if self.ops_type == 0:
+            ops = [self.gnn_map(
+                op, self.hidden_dim, self.hidden_dim, self.dropout)for op in self.gnn_ops]
+        elif self.ops_type == 1:
+            ops = [self.gnn_map(op, self.hidden_dim, self.hidden_dim,
+                                self.head, self.dropout)for op in self.gnn_ops]
+        choice = self.setLayerChoice(
+            layer,
+            ops,
+            key=key,
+        )
+        setattr(self, key, choice)
+        return choice
+
+    def _set_input_choice(self, key, layer, **kwargs):
+        setattr(self,
+                key,
+                self.setInputChoice(
+                    layer,
+                    key=key,
+                    **kwargs
+                ))
+
+    def forward(self, data):
+        x = bk_feat(data)
+        x = F.dropout(x, p=self.dropout, training=self.training)
+        prev_ = self.preproc0(x)
+
+        side_outs = []
+        stem_outs = []
+        input = prev_  
+        for layer in range(1, self.layer_number + 1):
+            # do layer choice for stem
+            op = getattr(self, f"stem_{layer}")
+            stem_out = bk_gconv(op, data, input)
+            stem_out = self.act(stem_out)  
+
+            # do double layer choice for sides
+            side_out_list = []
+            for i in range(2):
+                op = getattr(self, f'side_{layer}_{i}')
+                side_out = bk_gconv(op, data, input)
+                side_out = self.act(side_out)  # torch.Size([2, 2708, 64])
+                side_out_list.append(side_out)
+            side_out = torch.stack(side_out_list, dim=0)
+
+            stem_outs.append(stem_out)
+            side_outs.append(side_out)
+
+            # select input [x1,x2,x3] from side1,side2,stem
+            side_selected = getattr(self, f"in_{layer}")(side_outs)
+            input = [stem_outs[-1], side_selected]
+
+            # do agg in [add , attn]
+            agg = getattr(self, f"agg_{layer}")
+            input = bk_gconv(agg, data, input)
+
+        x = self.classifier2(input)
+        return F.log_softmax(x, dim=1)
+
+    def parse_model(self, selection, device) -> BaseModel:
+        return self.wrap(device).fix(selection)

autogl/module/nas/space/autoattend_space/operations.py

@@ -0,0 +1,216 @@
+
+from torch_geometric.nn import MessagePassing
+from torch_geometric.utils import softmax
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch_geometric.nn import GCNConv, SAGEConv, GATConv, ARMAConv, ChebConv, GatedGraphConv, SGConv
+
+from typing import Union, Tuple, Optional
+from torch_geometric.typing import (OptPairTensor, Adj, Size, NoneType,
+                                    OptTensor)
+
+import torch
+from torch import Tensor
+import torch.nn.functional as F
+from torch.nn import Parameter, Linear
+from torch_sparse import SparseTensor, set_diag
+from torch_geometric.nn.conv import MessagePassing
+from torch_geometric.utils import remove_self_loops, add_self_loops, softmax
+
+from torch_geometric.nn.inits import glorot, zeros
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import torch.nn.functional as F
+from torch.nn import Parameter
+from torch_geometric.nn.inits import glorot, zeros
+from torch_geometric.utils import softmax
+from torch_scatter import scatter_add
+import numpy as np
+from ..graph_nas_macro import GeoLayer
+
+
+class AggAdd(nn.Module):
+    def __init__(self, dim, att_head, dropout=0, norm=False, skip_connect=False, *args, **kwargs):
+        super(AggAdd, self).__init__()
+        self.dropout = dropout
+        self.ln_add = nn.BatchNorm1d(
+            dim, track_running_stats=True, affine=True)
+        self.norm = norm
+        self.skip_connect = skip_connect
+
+    def forward(self, x, edge_index, *args, **kwargs):
+        # x=[x_stem,[x_sides]]
+        norm = self.norm
+        x1, x2, x3 = x[0], x[1][0], x[1][1]
+        if norm:
+            return self.ln_add(x1 + x2)
+        else:
+            return x1 + x2
+
+
+class AggAttn(MessagePassing):
+    def __init__(self, dim, att_head, dropout=0, norm=False, skip_connect=False, *args, **kwargs):
+        super(AggAttn, self).__init__()
+        self.dropout = dropout
+        self.att_head = att_head
+        self.ln_attn = nn.BatchNorm1d(
+            dim, track_running_stats=True, affine=True)
+        self.norm = norm
+        self.skip_connect = skip_connect
+
+    def __repr__(self) -> str:
+        return 'AggAttn(att_head={}, dropout={})'.format(self.att_head, self.dropout)
+
+    def forward(self, x, edge_index, *args, **kwargs):
+        # x=[x_stem,[x_sides]]
+        # use dot-product attn
+        x1, x2, x3 = x[0], x[1][0], x[1][1]  # q,k,v
+        skip_connect, norm = self.skip_connect, self.norm
+        if not skip_connect and not norm:
+            return self.propagate(edge_index, x1=x1, x2=x2, x3=x3)
+
+        x = self.propagate(edge_index, x1=x1, x2=x2, x3=x3)
+        if not norm:
+            return x
+
+        if not skip_connect:
+            return self.ln_attn(x)
+        return self.ln_attn(x + x1)
+
+    def message(self, x2_j, x1_i, x3_j, index, ptr):
+        # x1: query, x2: key, x3: value # torch.Size([10556, 64]) ,index torch.Size([10556])
+        node, dim = x1_i.size()
+        dim_att = dim // self.att_head
+        # torch.Size([10556, 8, 8])
+        x2_j = x2_j.view(node, self.att_head, dim_att)
+        # torch.Size([10556, 8, 8])
+        x1_i = x1_i.view(node, self.att_head, dim_att)
+        attn = (x2_j * x1_i).sum(dim=-1) / \
+            np.sqrt(dim_att)  # torch.Size([10556, 8])
+        attn = softmax(attn, index, ptr)  # torch.Size([10556, 8])
+        # torch.Size([10556, 8])
+        attn = F.dropout(attn, p=self.dropout, training=self.training)
+        out = x3_j.view(node, self.att_head, dim_att) * attn.unsqueeze(-1)
+        out = out.view(-1, dim)
+        return out
+
+
+class GATConv2(MessagePassing):
+    _alpha: OptTensor
+
+    def __init__(self, in_channels: Union[int, Tuple[int, int]],
+                 out_channels: int, heads: int = 1, concat: bool = False,
+                 negative_slope: float = 0.2, dropout: float = 0.,
+                 add_self_loops: bool = True, bias: bool = True, **kwargs):
+        super(GATConv, self).__init__(aggr='add', node_dim=0, **kwargs)
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.heads = heads
+        self.concat = concat
+        self.negative_slope = negative_slope
+        self.dropout = dropout
+        self.add_self_loops = add_self_loops
+
+        self.lin = Linear(in_channels, heads * out_channels, bias=False)
+
+        self.att = Parameter(torch.Tensor(1, heads, out_channels))
+
+        if bias and concat:
+            self.bias = Parameter(torch.Tensor(heads * out_channels))
+        elif bias and not concat:
+            self.bias = Parameter(torch.Tensor(out_channels))
+        else:
+            self.register_parameter('bias', None)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        glorot(self.lin.weight)
+        glorot(self.att)
+        zeros(self.bias)
+
+    def forward(self, x: Union[Tensor, OptPairTensor], edge_index: Adj):
+        H, C = self.heads, self.out_channels
+
+        x = self.lin(x).view(-1, H, C)
+        alpha = (x * self.att).sum(dim=-1)
+
+        if self.add_self_loops:
+            if isinstance(edge_index, Tensor):
+                num_nodes = x.size(0)
+                edge_index, _ = remove_self_loops(edge_index)
+                edge_index, _ = add_self_loops(edge_index, num_nodes=num_nodes)
+            elif isinstance(edge_index, SparseTensor):
+                edge_index = set_diag(edge_index)
+
+        out = self.propagate(edge_index, x=x,
+                             alpha=alpha)
+
+        if self.concat:
+            out = out.view(-1, self.heads * self.out_channels)
+        else:
+            out = out.mean(dim=1)
+
+        if self.bias is not None:
+            out += self.bias
+
+        return out
+
+    def message(self, x_j: Tensor, alpha_j: Tensor, alpha_i: OptTensor,
+                index: Tensor, ptr: OptTensor) -> Tensor:
+        alpha = alpha_j if alpha_i is None else alpha_j + alpha_i
+        alpha = F.leaky_relu(alpha, self.negative_slope)
+        alpha = softmax(alpha, index, ptr)
+        alpha = F.dropout(alpha, p=self.dropout, training=self.training)
+        return x_j * alpha.unsqueeze(-1)
+
+    def __repr__(self):
+        return '{}({}, {}, heads={})'.format(self.__class__.__name__,
+                                             self.in_channels,
+                                             self.out_channels, self.heads)
+
+
+class Zero(nn.Module):
+    def __init__(self, indim, outdim) -> None:
+        super().__init__()
+        self.outdim = outdim
+        self.zero = nn.Parameter(torch.tensor(0.), requires_grad=True)
+
+    def forward(self, x, edge_index):
+        return torch.zeros(x.size(0), self.outdim).to(x.device) * self.zero
+
+# class Zero(nn.Module):
+#     def __init__(self, indim, outdim) -> None:
+#         super().__init__()
+#         self.outdim = outdim
+#         self.ln = nn.Linear(1, 1)
+
+#     def forward(self, x, edge_index):
+#         return 0.
+
+
+class Identity(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def forward(self, x, edge_index):
+        return x
+
+
+class Linear(nn.Module):
+    def __init__(self, in_dim, out_dim):
+        super().__init__()
+        self.core = nn.Linear(in_dim, out_dim)
+
+    def forward(self, x, *args):
+        return self.core(x)
+
+
+agg_map = {
+    'add': lambda dim, att_head=None, dropout=0, norm=False, skip_connect=False: AggAdd(dim, att_head, dropout, norm, skip_connect),
+    'attn': lambda dim, att_head=None, dropout=0, norm=False, skip_connect=False: AggAttn(dim, att_head, dropout, norm, skip_connect),
+}

autogl/module/nas/space/autoattend_space/ops1.py

@@ -0,0 +1,19 @@
+from .operations import *
+
+OPS = {
+    'ZERO': lambda indim, outdim, dropout, concat=False: Zero(indim, outdim),
+    'IDEN': lambda indim, outdim, dropout, concat=False: Identity(),
+    'GCN': lambda indim, outdim, dropout, concat=False: GCNConv(indim, outdim, add_self_loops=False),
+    'SAGE-MEAN': lambda indim, outdim, dropout, concat=False: SAGEConv(indim, outdim),
+    'GAT16': lambda indim, outdim, dropout, concat=False: GATConv(indim, outdim, dropout=dropout, heads=16, concat=False, add_self_loops=False) if not concat else GATConv(indim, outdim // 16, dropout=dropout, heads=16, concat=True, add_self_loops=False),
+    'GAT2': lambda indim, outdim, dropout, concat=False: GATConv(indim, outdim, dropout=dropout, heads=2, concat=False, add_self_loops=False) if not concat else GATConv(indim, outdim // 2, dropout=dropout, heads=2, concat=True, add_self_loops=False),
+    'GAT4': lambda indim, outdim, dropout, concat=False: GATConv(indim, outdim, dropout=dropout, heads=4, concat=False, add_self_loops=False) if not concat else GATConv(indim, outdim // 4, dropout=dropout, heads=4, concat=True, add_self_loops=False),
+    'GAT8': lambda indim, outdim, dropout, concat=False: GATConv(indim, outdim, dropout=dropout, heads=8, concat=False, add_self_loops=False) if not concat else GATConv(indim, outdim // 8, dropout=dropout, heads=8, concat=True, add_self_loops=False),
+    'GAT1': lambda indim, outdim, dropout, concat=False: GATConv(indim, outdim, dropout=dropout, heads=1, concat=False, add_self_loops=False),
+    'LIN': lambda indim, outdim, dropout, concat=False: Linear(indim, outdim),
+    'ARMA': lambda indim, outdim, dropout, concat=False: ARMAConv(indim, outdim),
+    'CHEB': lambda indim, outdim, dropout, concat=False: ChebConv(indim, outdim, 2),
+    'SGC': lambda indim, outdim, dropout, concat=False: SGConv(indim, outdim, add_self_loops=False)
+}
+
+PRIMITIVES = list(OPS.keys())

autogl/module/nas/space/autoattend_space/ops2.py

@@ -0,0 +1,13 @@
+from .operations import *
+OPS = {
+    'ZERO': lambda indim, outdim, head, dropout, concat=False: Zero(indim, outdim),
+    'CONST': lambda indim, outdim, head, dropout, concat=False: GeoLayer(indim, outdim, head, concat, att_type='const', dropout=dropout),
+    'GCN': lambda indim, outdim, head, dropout, concat=False: GeoLayer(indim, outdim, head, concat, att_type='gcn', dropout=dropout),
+    'GAT': lambda indim, outdim, head, dropout, concat=False: GeoLayer(indim, outdim, head, concat, att_type='gat', dropout=dropout),
+    'SYM': lambda indim, outdim, head, dropout, concat=False: GeoLayer(indim, outdim, head, concat, att_type='gat_sym', dropout=dropout),
+    'COS': lambda indim, outdim, head, dropout, concat=False: GeoLayer(indim, outdim, head, concat, att_type='cos', dropout=dropout),
+    'LIN': lambda indim, outdim, head, dropout, concat=False: GeoLayer(indim, outdim, head, concat, att_type='linear', dropout=dropout),
+    'GENE': lambda indim, outdim, head, dropout, concat=False: GeoLayer(indim, outdim, head, concat, att_type='generalized_linear', dropout=dropout)
+}
+
+PRIMITIVES = list(OPS.keys())

autogl/module/nas/space/base.py

@@ -10,7 +10,6 @@ from ....utils import get_logger
 from ..utils import get_hardware_aware_metric
 
 
-
 class OrderedMutable:
     """
     An abstract class with order, enabling to sort mutables with a certain rank.
@@ -30,7 +29,8 @@ class OrderedLayerChoice(OrderedMutable, mutables.LayerChoice):
         self, order, op_candidates, reduction="sum", return_mask=False, key=None
     ):
         OrderedMutable.__init__(self, order)
-        mutables.LayerChoice.__init__(self, op_candidates, reduction, return_mask, key)
+        mutables.LayerChoice.__init__(
+            self, op_candidates, reduction, return_mask, key)
 
 
 class OrderedInputChoice(OrderedMutable, mutables.InputChoice):
@@ -98,16 +98,15 @@ class BoxModel(BaseAutoModel):
 
     _logger = get_logger("space model")
 
-    def __init__(self, space_model, device=torch.device("cuda")):
+    def __init__(self, space_model, device):
         super().__init__(None, None, device)
         self.init = True
         self.space = []
         self.hyperparams = {}
-        self._model = space_model.to(device)
+        self._model = space_model
         self.num_features = self._model.input_dim
         self.num_classes = self._model.output_dim
         self.params = {"num_class": self.num_classes, "features_num": self.num_features}
-        self.device = device
         self.selection = None
 
     def _initialize(self):
@@ -139,11 +138,14 @@ class BoxModel(BaseAutoModel):
         ret_self._model.instantiate()
         if ret_self.selection:
             apply_fixed_architecture(ret_self._model, ret_self.selection, verbose=False)
-        ret_self.to_device(self.device)
         return ret_self
 
     def __repr__(self) -> str:
-        return str({'parameter': get_hardware_aware_metric(self.model, 'parameter')})
+        return str(
+            {'parameter': get_hardware_aware_metric(self.model, 'parameter'),
+             'model': self.model,
+             'selection': self.selection
+             })
 
 class BaseSpace(nn.Module):
     """
@@ -214,7 +216,8 @@ class BaseSpace(nn.Module):
             key = f"default_key_{self._default_key}"
             self._default_key += 1
             orikey = key
-        layer = OrderedLayerChoice(order, op_candidates, reduction, return_mask, orikey)
+        layer = OrderedLayerChoice(
+            order, op_candidates, reduction, return_mask, orikey)
         return layer
 
     def setInputChoice(
@@ -240,12 +243,13 @@ class BaseSpace(nn.Module):
         )
         return layer
 
-    def wrap(self, device="cuda"):
+    def wrap(self):
         """
         Return a BoxModel which wrap self as a model
         Used to pass to trainer
         To use this function, must contain `input_dim` and `output_dim`
         """
+        device = next(self.parameters()).device
         return BoxModel(self, device)
 
 
@@ -304,6 +308,7 @@ class CleanFixedArchitecture(FixedArchitecture):
         prefix : str
             Module name under global namespace.
         """
+        
         if module is None:
             module = self.model
         for name, mutable in module.named_children():

autogl/module/nas/space/gasso.py

@@ -0,0 +1,281 @@
+import typing as _typ
+
+from . import register_nas_space
+from .base import apply_fixed_architecture
+from .base import BaseSpace
+from ...model import BaseAutoModel
+from ....utils import get_logger
+
+from ..backend import *
+from ..utils import count_parameters, measure_latency
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import Module
+from .gasso_space import *
+from torch.autograd import Variable
+from collections import namedtuple
+
+Genotype = namedtuple('Genotype', 'normal normal_concat')
+Genotype_normal = namedtuple('Genotype_normal', 'normal normal_concat')
+
+gnn_list = [
+    "gat",  # GAT with 2 heads
+    "gcn",  # GCN
+    "gin",  # GIN
+    #"cheb",  # chebnet
+    "sage",  # sage
+    #"arma",
+    #"sg",  # simplifying gcn
+    "linear",  # skip connection
+    #"skip",  # skip connection
+    #"zero",  # skip connection
+]
+act_list = [
+    "sigmoid", "tanh", "relu", "linear", "elu"
+]
+
+def gnn_map(gnn_name, in_dim, out_dim, concat=False, bias=True) -> Module:
+    '''
+
+    :param gnn_name:
+    :param in_dim:
+    :param out_dim:
+    :param concat: for gat, concat multi-head output or not
+    :return: GNN model
+    '''
+    norm= True
+    if gnn_name == "gat":
+        return GATConv(in_dim, out_dim, 1, bias=bias, concat = False, add_self_loops=norm)
+    elif gnn_name == "gcn":
+        return GCNConv(in_dim, out_dim, add_self_loops=True, normalize=norm)
+    elif gnn_name == "gin":
+        return GINConv(torch.nn.Linear(in_dim, out_dim))
+    elif gnn_name == "cheb":
+        return ChebConv(in_dim, out_dim, K=2, bias=bias)
+    elif gnn_name == "sage":
+        return SAGEConv(in_dim, out_dim, bias=bias)
+    elif gnn_name == "gated":
+        return GatedGraphConv(in_dim, out_dim, bias=bias)
+    elif gnn_name == "arma":
+        return ARMAConv(in_dim, out_dim, bias=bias, normalize=norm)
+    elif gnn_name == "sg":
+        return SGConv(in_dim, out_dim, bias=bias, normalize=norm)
+    elif gnn_name == "linear":
+        return LinearConv(in_dim, out_dim, bias=bias)
+    elif gnn_name == "skip":
+        return SkipConv(in_dim, out_dim, bias=bias)
+    elif gnn_name == "zero":
+        return ZeroConv(in_dim, out_dim, bias=bias)
+    else:
+        raise ValueError("No such GNN name") 
+
+def Get_edges(adjs, ):
+    edges = []
+    edges_weights = []
+    for adj in adjs: 
+        edges.append(adj[0])
+        edges_weights.append(torch.sigmoid(adj[1]))
+    return edges, edges_weights
+
+class LinearConv(Module):
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 bias=True):
+        super(LinearConv, self).__init__()
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.linear = torch.nn.Linear(in_channels, out_channels, bias)
+
+    def forward(self, x, edge_index, edge_weight=None):
+        return self.linear(x)
+
+    def __repr__(self):
+        return '{}({}, {})'.format(self.__class__.__name__, self.in_channels,
+                                   self.out_channels)
+
+class SkipConv(Module):
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 bias=True):
+        super(SkipConv, self).__init__()
+        self.out_dim = out_channels
+
+
+    def forward(self, x, edge_index, edge_weight=None):
+        return x
+
+    def __repr__(self):
+        return '{}({}, {})'.format(self.__class__.__name__, self.in_channels,
+                                   self.out_channels)
+
+class ZeroConv(Module):
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 bias=True):
+        super(ZeroConv, self).__init__()
+        self.out_dim = out_channels
+
+
+    def forward(self, x, edge_index, edge_weight=None):
+        return torch.zeros([x.size(0), self.out_dim]).to(x.device)
+
+    def __repr__(self):
+        return '{}({}, {})'.format(self.__class__.__name__, self.in_channels,
+                                   self.out_channels)
+
+class MixedOp(nn.Module):
+
+    def __init__(self, in_c, out_c):
+        super(MixedOp, self).__init__()
+        self._ops = nn.ModuleList()
+        for action in gnn_list:
+            self._ops.append(gnn_map(action, in_c, out_c))
+
+    def forward(self, x, edge_index, edge_weight, weights, selected_idx=None):
+        if selected_idx is None:
+            fin = []
+            for w, op, op_name in zip(weights, self._ops, gnn_list):
+                """if op_name == "gcn":
+                    w = 1.0
+                else:
+                    continue"""
+                if edge_weight == None:
+                    fin.append(w * op(x, edge_index))
+                else:
+                    fin.append(w * op(x, edge_index, edge_weight = edge_weight))
+            return sum(fin)
+            #return sum(w * op(x, edge_index) for w, op in zip(weights, self._ops))
+        else:  # unchosen operations are pruned
+            return self._ops[selected_idx](x, edge_index)
+
+class CellWS(nn.Module):
+
+    def __init__(self, steps, his_dim, hidden_dim, out_dim, dp, bias=True):
+        super(CellWS, self).__init__()
+        self.steps = steps
+        self._ops = nn.ModuleList()
+        self._bns = nn.ModuleList()
+        self.use2 = False
+        self.dp = 0.8
+        for i in range(self.steps):
+            if i == 0: 
+                inpdim = his_dim
+            else:
+                inpdim = hidden_dim
+            if i == self.steps - 1:
+                oupdim = out_dim
+            else:
+                oupdim = hidden_dim
+            op = MixedOp(inpdim, oupdim)
+            self._ops.append(op)
+            self._bns.append(nn.BatchNorm1d(oupdim))
+
+    def forward(self, x, adjs, weights):
+        edges, ews = Get_edges(adjs)
+        for i in range(self.steps):
+            if i > 0:
+                x = F.relu(x)
+                x = F.dropout(x, p=self.dp, training=self.training)
+            x = self._ops[i](x, edges[i], ews[i], weights[i])  # call the gcn module
+        return x
+
+@register_nas_space("gassospace")
+class GassoSpace(BaseSpace):
+    def __init__(
+        self,
+        hidden_dim: _typ.Optional[int] = 64,
+        layer_number: _typ.Optional[int] = 2,
+        dropout: _typ.Optional[float] = 0.8,
+        input_dim: _typ.Optional[int] = None,
+        output_dim: _typ.Optional[int] = None,
+        ops: _typ.Tuple = gnn_list,
+    ):
+        super().__init__()
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+        self.hidden_dim = hidden_dim
+        self.steps = layer_number
+        self.dropout = dropout
+        self.ops = ops
+        self.use_forward = True
+        self.dead_tensor = torch.nn.Parameter(torch.FloatTensor([1]), requires_grad = True)
+
+    def instantiate(
+        self,
+        hidden_dim: _typ.Optional[int] = 64,
+        layer_number: _typ.Optional[int] = 2,
+        dropout: _typ.Optional[float] = 0.8,
+        input_dim: _typ.Optional[int] = None,
+        output_dim: _typ.Optional[int] = None,
+        ops: _typ.Tuple = gnn_list,
+    ):
+        super().instantiate()
+        self.input_dim = input_dim or self.input_dim
+        self.output_dim = output_dim or self.output_dim
+        self.hidden_dim = hidden_dim or self.hidden_dim
+        self.steps = layer_number or self.steps
+        self.dropout = dropout or self.dropout
+        self.ops = ops or self.ops
+        his_dim, cur_dim, hidden_dim, out_dim = self.input_dim, self.input_dim, self.hidden_dim, self.hidden_dim
+        self.cells = nn.ModuleList()
+
+        self.cell = CellWS(self.steps, his_dim, hidden_dim, self.output_dim, self.dropout)
+        his_dim = cur_dim
+        cur_dim = self.steps * out_dim
+
+        self.classifier = nn.Linear(cur_dim, self.output_dim)
+
+        self.initialize_alphas()
+
+    #def forward(self, x, adjs):
+    def forward(self, data):
+        if self.use_forward:
+            x, adjs = data.x, data.adj
+            x = F.dropout(x, p=self.dropout, training=self.training)
+
+            weights = []
+            for j in range(self.steps):
+                weights.append(F.softmax(self.alphas_normal[j], dim=-1))
+
+            x = self.cell(x, adjs, weights)
+            x = F.log_softmax(x, dim=1)
+            self.current_pred = x.detach()
+            return x
+        else:
+            #for i in self.parameters():
+            #    print(i)
+            x = self.prediction + self.dead_tensor * 0
+            return x
+
+    def keep_prediction(self):
+        self.prediction = self.current_pred
+
+    '''def to(self, *args, **kwargs):
+        fin = super().to(*args, **kwargs)
+        device = next(fin.parameters()).device
+        fin.alphas_normal = [i.to(device) for i in self.alphas_normal]
+        return fin'''
+
+    def initialize_alphas(self):
+        num_ops = len(self.ops)
+
+        self.alphas_normal = []
+        for i in range(self.steps):
+            self.alphas_normal.append(Variable(1e-3 * torch.randn(num_ops), requires_grad=True))
+
+        self._arch_parameters = [
+            self.alphas_normal
+        ]
+
+    def arch_parameters(self):
+        return self.alphas_normal
+
+    def parse_model(self, selection, device) -> BaseAutoModel:
+        self.use_forward = False
+        return self.wrap()

autogl/module/nas/space/gasso_space/__init__.py

@@ -0,0 +1,23 @@
+from .message_passing import MessagePassing
+from .gcn_conv import GCNConv
+from .cheb_conv import ChebConv
+from .sage_conv import SAGEConv
+from .gat_conv import GATConv
+from .gin_conv import GINConv, GINEConv
+from .arma_conv import ARMAConv
+from .edge_conv import EdgeConv, DynamicEdgeConv
+
+__all__ = [
+    'MessagePassing',
+    'GCNConv',
+    'ChebConv',
+    'SAGEConv',
+    'GATConv',
+    'GINConv',
+    'GINEConv',
+    'ARMAConv',
+    'EdgeConv',
+    'DynamicEdgeConv',
+]
+
+classes = __all__

autogl/module/nas/space/gasso_space/arma_conv.py

@@ -0,0 +1,133 @@
+from typing import Callable
+from torch_geometric.typing import Adj, OptTensor
+
+import torch
+from torch import Tensor
+from torch.nn import Parameter
+import torch.nn.functional as F
+from torch_sparse import SparseTensor, matmul
+from torch_geometric.nn.conv import MessagePassing
+from torch_geometric.nn.conv.gcn_conv import gcn_norm
+
+from .inits import glorot, zeros
+
+
+class ARMAConv(MessagePassing):
+    r"""The ARMA graph convolutional operator from the `"Graph Neural Networks
+    with Convolutional ARMA Filters" <https://arxiv.org/abs/1901.01343>`_ paper
+
+    .. math::
+        \mathbf{X}^{\prime} = \frac{1}{K} \sum_{k=1}^K \mathbf{X}_k^{(T)},
+
+    with :math:`\mathbf{X}_k^{(T)}` being recursively defined by
+
+    .. math::
+        \mathbf{X}_k^{(t+1)} = \sigma \left( \mathbf{\hat{L}}
+        \mathbf{X}_k^{(t)} \mathbf{W} + \mathbf{X}^{(0)} \mathbf{V} \right),
+
+    where :math:`\mathbf{\hat{L}} = \mathbf{I} - \mathbf{L} = \mathbf{D}^{-1/2}
+    \mathbf{A} \mathbf{D}^{-1/2}` denotes the
+    modified Laplacian :math:`\mathbf{L} = \mathbf{I} - \mathbf{D}^{-1/2}
+    \mathbf{A} \mathbf{D}^{-1/2}`.
+
+    Args:
+        in_channels (int): Size of each input sample :math:`\mathbf{x}^{(t)}`.
+        out_channels (int): Size of each output sample
+            :math:`\mathbf{x}^{(t+1)}`.
+        num_stacks (int, optional): Number of parallel stacks :math:`K`.
+            (default: :obj:`1`).
+        num_layers (int, optional): Number of layers :math:`T`.
+            (default: :obj:`1`)
+        act (callable, optional): Activation function :math:`\sigma`.
+            (default: :meth:`torch.nn.functional.ReLU`)
+        shared_weights (int, optional): If set to :obj:`True` the layers in
+            each stack will share the same parameters. (default: :obj:`False`)
+        dropout (float, optional): Dropout probability of the skip connection.
+            (default: :obj:`0.`)
+        bias (bool, optional): If set to :obj:`False`, the layer will not learn
+            an additive bias. (default: :obj:`True`)
+        **kwargs (optional): Additional arguments of
+            :class:`torch_geometric.nn.conv.MessagePassing`.
+    """
+    def __init__(self, in_channels: int, out_channels: int,
+                 num_stacks: int = 1, num_layers: int = 1,
+                 shared_weights: bool = False, act: Callable = F.relu, normalize = True,
+                 dropout: float = 0., bias: bool = True, **kwargs):
+        kwargs.setdefault('aggr', 'add')
+        super(ARMAConv, self).__init__(**kwargs)
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.num_stacks = num_stacks
+        self.num_layers = num_layers
+        self.act = act
+        self.shared_weights = shared_weights
+        self.dropout = dropout
+        self.normalize = normalize
+
+        K, T, F_in, F_out = num_stacks, num_layers, in_channels, out_channels
+        T = 1 if self.shared_weights else T
+
+        self.init_weight = Parameter(torch.Tensor(K, F_in, F_out))
+        self.weight = Parameter(torch.Tensor(max(1, T - 1), K, F_out, F_out))
+        self.root_weight = Parameter(torch.Tensor(T, K, F_in, F_out))
+
+        if bias:
+            self.bias = Parameter(torch.Tensor(T, K, 1, F_out))
+        else:
+            self.register_parameter('bias', None)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        glorot(self.init_weight)
+        glorot(self.weight)
+        glorot(self.root_weight)
+        zeros(self.bias)
+
+    def forward(self, x: Tensor, edge_index: Adj,
+                edge_weight: OptTensor = None) -> Tensor:
+        """"""
+
+        if isinstance(edge_index, Tensor) and self.normalize:
+            edge_index, edge_weight = gcn_norm(  # yapf: disable
+                edge_index, edge_weight, x.size(self.node_dim),
+                add_self_loops=False, dtype=x.dtype)
+
+        elif isinstance(edge_index, SparseTensor) and self.normalize:
+            edge_index = gcn_norm(  # yapf: disable
+                edge_index, edge_weight, x.size(self.node_dim),
+                add_self_loops=False, dtype=x.dtype)
+
+        x = x.unsqueeze(-3)
+        out = x
+        for t in range(self.num_layers):
+            if t == 0:
+                out = out @ self.init_weight
+            else:
+                out = out @ self.weight[0 if self.shared_weights else t - 1]
+
+            # propagate_type: (x: Tensor, edge_weight: OptTensor)
+            out = self.propagate(edge_index, x=out, edge_weight=edge_weight,
+                                 size=None)
+
+            root = F.dropout(x, p=self.dropout, training=self.training)
+            out += root @ self.root_weight[0 if self.shared_weights else t]
+
+            if self.bias is not None:
+                out += self.bias[0 if self.shared_weights else t]
+
+            out = self.act(out)
+
+        return out.mean(dim=-3)
+
+    def message(self, x_j: Tensor, edge_weight: Tensor) -> Tensor:
+        return edge_weight.view(-1, 1) * x_j
+
+    def message_and_aggregate(self, adj_t: SparseTensor, x: Tensor) -> Tensor:
+        return matmul(adj_t, x, reduce=self.aggr)
+
+    def __repr__(self):
+        return '{}({}, {}, num_stacks={}, num_layers={})'.format(
+            self.__class__.__name__, self.in_channels, self.out_channels,
+            self.num_stacks, self.num_layers)

autogl/module/nas/space/gasso_space/cheb_conv.py

@@ -0,0 +1,157 @@
+from typing import Optional
+from torch_geometric.typing import OptTensor
+
+import torch
+from torch.nn import Parameter
+from torch_geometric.nn.conv import MessagePassing
+from torch_geometric.utils import remove_self_loops, add_self_loops
+from torch_geometric.utils import get_laplacian
+
+from .inits import glorot, zeros
+
+
+class ChebConv(MessagePassing):
+    r"""The chebyshev spectral graph convolutional operator from the
+    `"Convolutional Neural Networks on Graphs with Fast Localized Spectral
+    Filtering" <https://arxiv.org/abs/1606.09375>`_ paper
+
+    .. math::
+        \mathbf{X}^{\prime} = \sum_{k=1}^{K} \mathbf{Z}^{(k)} \cdot
+        \mathbf{\Theta}^{(k)}
+
+    where :math:`\mathbf{Z}^{(k)}` is computed recursively by
+
+    .. math::
+        \mathbf{Z}^{(1)} &= \mathbf{X}
+
+        \mathbf{Z}^{(2)} &= \mathbf{\hat{L}} \cdot \mathbf{X}
+
+        \mathbf{Z}^{(k)} &= 2 \cdot \mathbf{\hat{L}} \cdot
+        \mathbf{Z}^{(k-1)} - \mathbf{Z}^{(k-2)}
+
+    and :math:`\mathbf{\hat{L}}` denotes the scaled and normalized Laplacian
+    :math:`\frac{2\mathbf{L}}{\lambda_{\max}} - \mathbf{I}`.
+
+    Args:
+        in_channels (int): Size of each input sample.
+        out_channels (int): Size of each output sample.
+        K (int): Chebyshev filter size :math:`K`.
+        normalization (str, optional): The normalization scheme for the graph
+            Laplacian (default: :obj:`"sym"`):
+
+            1. :obj:`None`: No normalization
+            :math:`\mathbf{L} = \mathbf{D} - \mathbf{A}`
+
+            2. :obj:`"sym"`: Symmetric normalization
+            :math:`\mathbf{L} = \mathbf{I} - \mathbf{D}^{-1/2} \mathbf{A}
+            \mathbf{D}^{-1/2}`
+
+            3. :obj:`"rw"`: Random-walk normalization
+            :math:`\mathbf{L} = \mathbf{I} - \mathbf{D}^{-1} \mathbf{A}`
+
+            You need to pass :obj:`lambda_max` to the :meth:`forward` method of
+            this operator in case the normalization is non-symmetric.
+            :obj:`\lambda_max` should be a :class:`torch.Tensor` of size
+            :obj:`[num_graphs]` in a mini-batch scenario and a
+            scalar/zero-dimensional tensor when operating on single graphs.
+            You can pre-compute :obj:`lambda_max` via the
+            :class:`torch_geometric.transforms.LaplacianLambdaMax` transform.
+        bias (bool, optional): If set to :obj:`False`, the layer will not learn
+            an additive bias. (default: :obj:`True`)
+        **kwargs (optional): Additional arguments of
+            :class:`torch_geometric.nn.conv.MessagePassing`.
+    """
+    def __init__(self, in_channels, out_channels, K, normalization='sym',
+                 bias=True, **kwargs):
+        kwargs.setdefault('aggr', 'add')
+        super(ChebConv, self).__init__(**kwargs)
+
+        assert K > 0
+        assert normalization in [None, 'sym', 'rw'], 'Invalid normalization'
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.normalization = normalization
+        self.weight = Parameter(torch.Tensor(K, in_channels, out_channels))
+
+        if bias:
+            self.bias = Parameter(torch.Tensor(out_channels))
+        else:
+            self.register_parameter('bias', None)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        glorot(self.weight)
+        zeros(self.bias)
+
+    def __norm__(self, edge_index, num_nodes: Optional[int],
+                 edge_weight: OptTensor, normalization: Optional[str],
+                 lambda_max, dtype: Optional[int] = None,
+                 batch: OptTensor = None):
+
+        edge_index, edge_weight = remove_self_loops(edge_index, edge_weight)
+
+        edge_index, edge_weight = get_laplacian(edge_index, edge_weight,
+                                                normalization, dtype,
+                                                num_nodes)
+
+        if batch is not None and lambda_max.numel() > 1:
+            lambda_max = lambda_max[batch[edge_index[0]]]
+
+        edge_weight = (2.0 * edge_weight) / lambda_max
+        edge_weight.masked_fill_(edge_weight == float('inf'), 0)
+
+        edge_index, edge_weight = add_self_loops(edge_index, edge_weight,
+                                                 fill_value=-1.,
+                                                 num_nodes=num_nodes)
+        assert edge_weight is not None
+
+        return edge_index, edge_weight
+
+    def forward(self, x, edge_index, edge_weight: OptTensor = None,
+                batch: OptTensor = None, lambda_max: OptTensor = None):
+        """"""
+        if self.normalization != 'sym' and lambda_max is None:
+            raise ValueError('You need to pass `lambda_max` to `forward() in`'
+                             'case the normalization is non-symmetric.')
+
+        if lambda_max is None:
+            lambda_max = torch.tensor(2.0, dtype=x.dtype, device=x.device)
+        if not isinstance(lambda_max, torch.Tensor):
+            lambda_max = torch.tensor(lambda_max, dtype=x.dtype,
+                                      device=x.device)
+        assert lambda_max is not None
+
+        edge_index, norm = self.__norm__(edge_index, x.size(self.node_dim),
+                                         edge_weight, self.normalization,
+                                         lambda_max, dtype=x.dtype,
+                                         batch=batch)
+
+        Tx_0 = x
+        Tx_1 = x  # Dummy.
+        out = torch.matmul(Tx_0, self.weight[0])
+
+        # propagate_type: (x: Tensor, norm: Tensor)
+        if self.weight.size(0) > 1:
+            Tx_1 = self.propagate(edge_index, x=x, norm=norm, size=None)
+            out = out + torch.matmul(Tx_1, self.weight[1])
+
+        for k in range(2, self.weight.size(0)):
+            Tx_2 = self.propagate(edge_index, x=Tx_1, norm=norm, size=None)
+            Tx_2 = 2. * Tx_2 - Tx_0
+            out = out + torch.matmul(Tx_2, self.weight[k])
+            Tx_0, Tx_1 = Tx_1, Tx_2
+
+        if self.bias is not None:
+            out += self.bias
+
+        return out
+
+    def message(self, x_j, norm):
+        return norm.view(-1, 1) * x_j
+
+    def __repr__(self):
+        return '{}({}, {}, K={}, normalization={})'.format(
+            self.__class__.__name__, self.in_channels, self.out_channels,
+            self.weight.size(0), self.normalization)

autogl/module/nas/space/gasso_space/edge_conv.py

@@ -0,0 +1,123 @@
+from typing import Callable, Union, Optional
+from torch_geometric.typing import OptTensor, PairTensor, PairOptTensor, Adj
+
+import torch
+from torch import Tensor
+from torch_geometric.nn.conv import MessagePassing
+
+from .inits import reset
+
+try:
+    from torch_cluster import knn
+except ImportError:
+    knn = None
+
+
+class EdgeConv(MessagePassing):
+    r"""The edge convolutional operator from the `"Dynamic Graph CNN for
+    Learning on Point Clouds" <https://arxiv.org/abs/1801.07829>`_ paper
+
+    .. math::
+        \mathbf{x}^{\prime}_i = \sum_{j \in \mathcal{N}(i)}
+        h_{\mathbf{\Theta}}(\mathbf{x}_i \, \Vert \,
+        \mathbf{x}_j - \mathbf{x}_i),
+
+    where :math:`h_{\mathbf{\Theta}}` denotes a neural network, *.i.e.* a MLP.
+
+    Args:
+        nn (torch.nn.Module): A neural network :math:`h_{\mathbf{\Theta}}` that
+            maps pair-wise concatenated node features :obj:`x` of shape
+            :obj:`[-1, 2 * in_channels]` to shape :obj:`[-1, out_channels]`,
+            *e.g.*, defined by :class:`torch.nn.Sequential`.
+        aggr (string, optional): The aggregation scheme to use
+            (:obj:`"add"`, :obj:`"mean"`, :obj:`"max"`).
+            (default: :obj:`"max"`)
+        **kwargs (optional): Additional arguments of
+            :class:`torch_geometric.nn.conv.MessagePassing`.
+    """
+    def __init__(self, nn: Callable, aggr: str = 'max', **kwargs):
+        super(EdgeConv, self).__init__(aggr=aggr, **kwargs)
+        self.nn = nn
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        reset(self.nn)
+
+    def forward(self, x: Union[Tensor, PairTensor], edge_index: Adj) -> Tensor:
+        """"""
+        if isinstance(x, Tensor):
+            x: PairTensor = (x, x)
+        # propagate_type: (x: PairTensor)
+        return self.propagate(edge_index, x=x, size=None)
+
+    def message(self, x_i: Tensor, x_j: Tensor) -> Tensor:
+        return self.nn(torch.cat([x_i, x_j - x_i], dim=-1))
+
+    def __repr__(self):
+        return '{}(nn={})'.format(self.__class__.__name__, self.nn)
+
+
+class DynamicEdgeConv(MessagePassing):
+    r"""The dynamic edge convolutional operator from the `"Dynamic Graph CNN
+    for Learning on Point Clouds" <https://arxiv.org/abs/1801.07829>`_ paper
+    (see :class:`torch_geometric.nn.conv.EdgeConv`), where the graph is
+    dynamically constructed using nearest neighbors in the feature space.
+
+    Args:
+        nn (torch.nn.Module): A neural network :math:`h_{\mathbf{\Theta}}` that
+            maps pair-wise concatenated node features :obj:`x` of shape
+            `:obj:`[-1, 2 * in_channels]` to shape :obj:`[-1, out_channels]`,
+            *e.g.* defined by :class:`torch.nn.Sequential`.
+        k (int): Number of nearest neighbors.
+        aggr (string): The aggregation operator to use (:obj:`"add"`,
+            :obj:`"mean"`, :obj:`"max"`). (default: :obj:`"max"`)
+        num_workers (int): Number of workers to use for k-NN computation.
+            Has no effect in case :obj:`batch` is not :obj:`None`, or the input
+            lies on the GPU. (default: :obj:`1`)
+        **kwargs (optional): Additional arguments of
+            :class:`torch_geometric.nn.conv.MessagePassing`.
+    """
+    def __init__(self, nn: Callable, k: int, aggr: str = 'max',
+                 num_workers: int = 1, **kwargs):
+        super(DynamicEdgeConv,
+              self).__init__(aggr=aggr, flow='target_to_source', **kwargs)
+
+        if knn is None:
+            raise ImportError('`DynamicEdgeConv` requires `torch-cluster`.')
+
+        self.nn = nn
+        self.k = k
+        self.num_workers = num_workers
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        reset(self.nn)
+
+    def forward(
+            self, x: Union[Tensor, PairTensor],
+            batch: Union[OptTensor, Optional[PairTensor]] = None) -> Tensor:
+        """"""
+        if isinstance(x, Tensor):
+            x: PairTensor = (x, x)
+        assert x[0].dim() == 2, \
+            'Static graphs not supported in `DynamicEdgeConv`.'
+
+        b: PairOptTensor = (None, None)
+        if isinstance(batch, Tensor):
+            b = (batch, batch)
+        elif isinstance(batch, tuple):
+            assert batch is not None
+            b = (batch[0], batch[1])
+
+        edge_index = knn(x[0], x[1], self.k, b[0], b[1],
+                         num_workers=self.num_workers)
+
+        # propagate_type: (x: PairTensor)
+        return self.propagate(edge_index, x=x, size=None)
+
+    def message(self, x_i: Tensor, x_j: Tensor) -> Tensor:
+        return self.nn(torch.cat([x_i, x_j - x_i], dim=-1))
+
+    def __repr__(self):
+        return '{}(nn={}, k={})'.format(self.__class__.__name__, self.nn,
+                                        self.k)

autogl/module/nas/space/gasso_space/gat_conv.py

@@ -0,0 +1,194 @@
+from typing import Union, Tuple, Optional
+from torch_geometric.typing import (OptPairTensor, Adj, Size, NoneType,
+                                    OptTensor)
+
+import torch
+from torch import Tensor
+import torch.nn.functional as F
+from torch.nn import Parameter, Linear
+from torch_sparse import SparseTensor, set_diag
+from torch_geometric.nn.conv import MessagePassing
+from torch_geometric.utils import remove_self_loops, add_self_loops, softmax
+
+from .inits import glorot, zeros
+
+
+class GATConv(MessagePassing):
+    r"""The graph attentional operator from the `"Graph Attention Networks"
+    <https://arxiv.org/abs/1710.10903>`_ paper
+
+    .. math::
+        \mathbf{x}^{\prime}_i = \alpha_{i,i}\mathbf{\Theta}\mathbf{x}_{i} +
+        \sum_{j \in \mathcal{N}(i)} \alpha_{i,j}\mathbf{\Theta}\mathbf{x}_{j},
+
+    where the attention coefficients :math:`\alpha_{i,j}` are computed as
+
+    .. math::
+        \alpha_{i,j} =
+        \frac{
+        \exp\left(\mathrm{LeakyReLU}\left(\mathbf{a}^{\top}
+        [\mathbf{\Theta}\mathbf{x}_i \, \Vert \, \mathbf{\Theta}\mathbf{x}_j]
+        \right)\right)}
+        {\sum_{k \in \mathcal{N}(i) \cup \{ i \}}
+        \exp\left(\mathrm{LeakyReLU}\left(\mathbf{a}^{\top}
+        [\mathbf{\Theta}\mathbf{x}_i \, \Vert \, \mathbf{\Theta}\mathbf{x}_k]
+        \right)\right)}.
+
+    Args:
+        in_channels (int or tuple): Size of each input sample. A tuple
+            corresponds to the sizes of source and target dimensionalities.
+        out_channels (int): Size of each output sample.
+        heads (int, optional): Number of multi-head-attentions.
+            (default: :obj:`1`)
+        concat (bool, optional): If set to :obj:`False`, the multi-head
+            attentions are averaged instead of concatenated.
+            (default: :obj:`True`)
+        negative_slope (float, optional): LeakyReLU angle of the negative
+            slope. (default: :obj:`0.2`)
+        dropout (float, optional): Dropout probability of the normalized
+            attention coefficients which exposes each node to a stochastically
+            sampled neighborhood during training. (default: :obj:`0`)
+        add_self_loops (bool, optional): If set to :obj:`False`, will not add
+            self-loops to the input graph. (default: :obj:`True`)
+        bias (bool, optional): If set to :obj:`False`, the layer will not learn
+            an additive bias. (default: :obj:`True`)
+        **kwargs (optional): Additional arguments of
+            :class:`torch_geometric.nn.conv.MessagePassing`.
+    """
+    _alpha: OptTensor
+
+    def __init__(self, in_channels: Union[int, Tuple[int, int]],
+                 out_channels: int, heads: int = 1, concat: bool = True,
+                 negative_slope: float = 0.2, dropout: float = 0.,
+                 add_self_loops: bool = True, bias: bool = True, **kwargs):
+        kwargs.setdefault('aggr', 'add')
+        super(GATConv, self).__init__(node_dim=0, **kwargs)
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.heads = heads
+        self.concat = concat
+        self.negative_slope = negative_slope
+        self.dropout = dropout
+        self.add_self_loops = add_self_loops
+
+        if isinstance(in_channels, int):
+            self.lin_l = Linear(in_channels, heads * out_channels, bias=False)
+            self.lin_r = self.lin_l
+        else:
+            self.lin_l = Linear(in_channels[0], heads * out_channels, False)
+            self.lin_r = Linear(in_channels[1], heads * out_channels, False)
+
+        self.att_l = Parameter(torch.Tensor(1, heads, out_channels))
+        self.att_r = Parameter(torch.Tensor(1, heads, out_channels))
+
+        if bias and concat:
+            self.bias = Parameter(torch.Tensor(heads * out_channels))
+        elif bias and not concat:
+            self.bias = Parameter(torch.Tensor(out_channels))
+        else:
+            self.register_parameter('bias', None)
+
+        self._alpha = None
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        glorot(self.lin_l.weight)
+        glorot(self.lin_r.weight)
+        glorot(self.att_l)
+        glorot(self.att_r)
+        zeros(self.bias)
+
+    def forward(self, x: Union[Tensor, OptPairTensor], edge_index: Adj, edge_weight: OptTensor = None,
+                size: Size = None, return_attention_weights=None):
+        # type: (Union[Tensor, OptPairTensor], Tensor, Size, NoneType) -> Tensor  # noqa
+        # type: (Union[Tensor, OptPairTensor], SparseTensor, Size, NoneType) -> Tensor  # noqa
+        # type: (Union[Tensor, OptPairTensor], Tensor, Size, bool) -> Tuple[Tensor, Tuple[Tensor, Tensor]]  # noqa
+        # type: (Union[Tensor, OptPairTensor], SparseTensor, Size, bool) -> Tuple[Tensor, SparseTensor]  # noqa
+        r"""
+
+        Args:
+            return_attention_weights (bool, optional): If set to :obj:`True`,
+                will additionally return the tuple
+                :obj:`(edge_index, attention_weights)`, holding the computed
+                attention weights for each edge. (default: :obj:`None`)
+        """
+        H, C = self.heads, self.out_channels
+
+        x_l: OptTensor = None
+        x_r: OptTensor = None
+        alpha_l: OptTensor = None
+        alpha_r: OptTensor = None
+        if isinstance(x, Tensor):
+            assert x.dim() == 2, 'Static graphs not supported in `GATConv`.'
+            x_l = x_r = self.lin_l(x).view(-1, H, C)
+            alpha_l = (x_l * self.att_l).sum(dim=-1)
+            alpha_r = (x_r * self.att_r).sum(dim=-1)
+        else:
+            x_l, x_r = x[0], x[1]
+            assert x[0].dim() == 2, 'Static graphs not supported in `GATConv`.'
+            x_l = self.lin_l(x_l).view(-1, H, C)
+            alpha_l = (x_l * self.att_l).sum(dim=-1)
+            if x_r is not None:
+                x_r = self.lin_r(x_r).view(-1, H, C)
+                alpha_r = (x_r * self.att_r).sum(dim=-1)
+
+        assert x_l is not None
+        assert alpha_l is not None
+
+        if self.add_self_loops:
+            if isinstance(edge_index, Tensor):
+                num_nodes = x_l.size(0)
+                if x_r is not None:
+                    num_nodes = min(num_nodes, x_r.size(0))
+                if size is not None:
+                    num_nodes = min(size[0], size[1])
+                edge_index, edge_weight = remove_self_loops(edge_index, edge_attr=edge_weight)
+                if edge_weight != None:
+                    edge_index, edge_weight = add_self_loops(edge_index, edge_weight=edge_weight, num_nodes=num_nodes)
+                else:
+                    edge_index, _ = add_self_loops(edge_index, num_nodes=num_nodes)
+            elif isinstance(edge_index, SparseTensor):
+                edge_index = set_diag(edge_index)
+
+        # propagate_type: (x: OptPairTensor, alpha: OptPairTensor)
+        out = self.propagate(edge_index, x=(x_l, x_r),
+                             alpha=(alpha_l, alpha_r), edge_weight = edge_weight, size=size)
+
+        alpha = self._alpha
+        self._alpha = None
+
+        if self.concat:
+            out = out.view(-1, self.heads * self.out_channels)
+        else:
+            out = out.mean(dim=1)
+
+        if self.bias is not None:
+            out += self.bias
+
+        if isinstance(return_attention_weights, bool):
+            assert alpha is not None
+            if isinstance(edge_index, Tensor):
+                return out, (edge_index, alpha)
+            elif isinstance(edge_index, SparseTensor):
+                return out, edge_index.set_value(alpha, layout='coo')
+        else:
+            return out
+
+    def message(self, x_j: Tensor, alpha_j: Tensor, alpha_i: OptTensor,
+                index: Tensor, ptr: OptTensor,
+                size_i: Optional[int], edge_weight: OptTensor = None) -> Tensor:
+        alpha = alpha_j if alpha_i is None else alpha_j + alpha_i
+        alpha = F.leaky_relu(alpha, self.negative_slope)
+        if edge_weight != None:
+            alpha = alpha.mul(edge_weight.unsqueeze(1))
+        alpha = softmax(alpha, index, ptr, size_i)
+        self._alpha = alpha
+        alpha = F.dropout(alpha, p=self.dropout, training=self.training)
+        return x_j * alpha.unsqueeze(-1)
+
+    def __repr__(self):
+        return '{}({}, {}, heads={})'.format(self.__class__.__name__,
+                                             self.in_channels,
+                                             self.out_channels, self.heads)

autogl/module/nas/space/gasso_space/gcn_conv.py

@@ -0,0 +1,200 @@
+from typing import Optional, Tuple
+from torch_geometric.typing import Adj, OptTensor, PairTensor
+
+import torch
+from torch import Tensor
+from torch.nn import Parameter
+from torch_scatter import scatter_add
+from torch_sparse import SparseTensor, matmul, fill_diag, sum, mul
+from torch_geometric.nn.conv import MessagePassing
+from torch_geometric.utils import add_remaining_self_loops
+from torch_geometric.utils.num_nodes import maybe_num_nodes
+
+from .inits import glorot, zeros
+
+
+@torch.jit._overload
+def gcn_norm(edge_index, edge_weight=None, num_nodes=None, improved=False,
+             add_self_loops=True, dtype=None):
+    # type: (Tensor, OptTensor, Optional[int], bool, bool, Optional[int]) -> PairTensor  # noqa
+    pass
+
+
+@torch.jit._overload
+def gcn_norm(edge_index, edge_weight=None, num_nodes=None, improved=False,
+             add_self_loops=True, dtype=None):
+    # type: (SparseTensor, OptTensor, Optional[int], bool, bool, Optional[int]) -> SparseTensor  # noqa
+    pass
+
+
+def gcn_norm(edge_index, edge_weight=None, num_nodes=None, improved=False,
+             add_self_loops=True, dtype=None):
+
+    fill_value = 2. if improved else 1.
+
+    if isinstance(edge_index, SparseTensor):
+        adj_t = edge_index
+        if not adj_t.has_value():
+            adj_t = adj_t.fill_value(1., dtype=dtype)
+        if add_self_loops:
+            adj_t = fill_diag(adj_t, fill_value)
+        deg = sum(adj_t, dim=1)
+        deg_inv_sqrt = deg.pow_(-0.5)
+        deg_inv_sqrt.masked_fill_(deg_inv_sqrt == float('inf'), 0.)
+        adj_t = mul(adj_t, deg_inv_sqrt.view(-1, 1))
+        adj_t = mul(adj_t, deg_inv_sqrt.view(1, -1))
+        return adj_t
+
+    else:
+        num_nodes = maybe_num_nodes(edge_index, num_nodes)
+
+        if edge_weight is None:
+            edge_weight = torch.ones((edge_index.size(1), ), dtype=dtype,
+                                     device=edge_index.device)
+
+        if add_self_loops:
+            edge_index, tmp_edge_weight = add_remaining_self_loops(
+                edge_index, edge_weight, fill_value, num_nodes)
+            assert tmp_edge_weight is not None
+            edge_weight = tmp_edge_weight
+
+        row, col = edge_index[0], edge_index[1]
+        deg = scatter_add(edge_weight, col, dim=0, dim_size=num_nodes)
+        deg_inv_sqrt = deg.pow_(-0.5)
+        deg_inv_sqrt.masked_fill_(deg_inv_sqrt == float('inf'), 0)
+        return edge_index, deg_inv_sqrt[row] * edge_weight * deg_inv_sqrt[col]
+
+
+class GCNConv(MessagePassing):
+    r"""The graph convolutional operator from the `"Semi-supervised
+    Classification with Graph Convolutional Networks"
+    <https://arxiv.org/abs/1609.02907>`_ paper
+
+    .. math::
+        \mathbf{X}^{\prime} = \mathbf{\hat{D}}^{-1/2} \mathbf{\hat{A}}
+        \mathbf{\hat{D}}^{-1/2} \mathbf{X} \mathbf{\Theta},
+
+    where :math:`\mathbf{\hat{A}} = \mathbf{A} + \mathbf{I}` denotes the
+    adjacency matrix with inserted self-loops and
+    :math:`\hat{D}_{ii} = \sum_{j=0} \hat{A}_{ij}` its diagonal degree matrix.
+    The adjacency matrix can include other values than :obj:`1` representing
+    edge weights via the optional :obj:`edge_weight` tensor.
+
+    Its node-wise formulation is given by:
+
+    .. math::
+        \mathbf{x}^{\prime}_i = \mathbf{\Theta} \sum_{j}
+        \frac{1}{\sqrt{\hat{d}_j \hat{d}_i}} \mathbf{x}_j
+
+    with :math:`\hat{d}_i = 1 + \sum_{j \in \mathcal{N}(i)} e_{j,i}`, where
+    :math:`e_{j,i}` denotes the edge weight from source node :obj:`i` to target
+    node :obj:`j` (default: :obj:`1`)
+
+    Args:
+        in_channels (int): Size of each input sample.
+        out_channels (int): Size of each output sample.
+        improved (bool, optional): If set to :obj:`True`, the layer computes
+            :math:`\mathbf{\hat{A}}` as :math:`\mathbf{A} + 2\mathbf{I}`.
+            (default: :obj:`False`)
+        cached (bool, optional): If set to :obj:`True`, the layer will cache
+            the computation of :math:`\mathbf{\hat{D}}^{-1/2} \mathbf{\hat{A}}
+            \mathbf{\hat{D}}^{-1/2}` on first execution, and will use the
+            cached version for further executions.
+            This parameter should only be set to :obj:`True` in transductive
+            learning scenarios. (default: :obj:`False`)
+        add_self_loops (bool, optional): If set to :obj:`False`, will not add
+            self-loops to the input graph. (default: :obj:`True`)
+        normalize (bool, optional): Whether to add self-loops and apply
+            symmetric normalization. (default: :obj:`True`)
+        bias (bool, optional): If set to :obj:`False`, the layer will not learn
+            an additive bias. (default: :obj:`True`)
+        **kwargs (optional): Additional arguments of
+            :class:`torch_geometric.nn.conv.MessagePassing`.
+    """
+
+    _cached_edge_index: Optional[Tuple[Tensor, Tensor]]
+    _cached_adj_t: Optional[SparseTensor]
+
+    def __init__(self, in_channels: int, out_channels: int,
+                 improved: bool = False, cached: bool = False,
+                 add_self_loops: bool = True, normalize: bool = True,
+                 bias: bool = True, **kwargs):
+
+        kwargs.setdefault('aggr', 'add')
+        super(GCNConv, self).__init__(**kwargs)
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.improved = improved
+        self.cached = cached
+        self.add_self_loops = add_self_loops
+        self.normalize = normalize
+
+        self._cached_edge_index = None
+        self._cached_adj_t = None
+
+        self.weight = Parameter(torch.Tensor(in_channels, out_channels))
+
+        if bias:
+            self.bias = Parameter(torch.Tensor(out_channels))
+        else:
+            self.register_parameter('bias', None)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        glorot(self.weight)
+        zeros(self.bias)
+        self._cached_edge_index = None
+        self._cached_adj_t = None
+
+    def forward(self, x: Tensor, edge_index: Adj,
+                edge_weight: OptTensor = None) -> Tensor:
+        """"""
+
+        if self.normalize:
+            if isinstance(edge_index, Tensor):
+                cache = self._cached_edge_index
+                if cache is None:
+                    edge_index, edge_weight = gcn_norm(  # yapf: disable
+                        edge_index, edge_weight, x.size(self.node_dim),
+                        self.improved, self.add_self_loops, dtype=x.dtype)
+                    if self.cached:
+                        self._cached_edge_index = (edge_index, edge_weight)
+                else:
+                    edge_index, edge_weight = cache[0], cache[1]
+
+            elif isinstance(edge_index, SparseTensor):
+                cache = self._cached_adj_t
+                if cache is None:
+                    edge_index = gcn_norm(  # yapf: disable
+                        edge_index, edge_weight, x.size(self.node_dim),
+                        self.improved, self.add_self_loops, dtype=x.dtype)
+                    if self.cached:
+                        self._cached_adj_t = edge_index
+                else:
+                    edge_index = cache
+
+        x = torch.matmul(x, self.weight)
+
+        # propagate_type: (x: Tensor, edge_weight: OptTensor)
+        out = self.propagate(edge_index, x=x, edge_weight=edge_weight,
+                             size=None)
+
+        if self.bias is not None:
+            out += self.bias
+
+        return out
+
+    def message(self, x_j: Tensor, edge_weight: OptTensor) -> Tensor:
+        if edge_weight is None:
+            return x_j
+        else:
+            return edge_weight.view(-1, 1) * x_j
+
+    def message_and_aggregate(self, adj_t: SparseTensor, x: Tensor) -> Tensor:
+        return matmul(adj_t, x, reduce=self.aggr)
+
+    def __repr__(self):
+        return '{}({}, {})'.format(self.__class__.__name__, self.in_channels,
+                                   self.out_channels)

autogl/module/nas/space/gasso_space/gin_conv.py

@@ -0,0 +1,157 @@
+from typing import Callable, Union
+from torch_geometric.typing import OptPairTensor, Adj, OptTensor, Size
+
+import torch
+from torch import Tensor
+import torch.nn.functional as F
+from torch_sparse import SparseTensor, matmul
+from torch_geometric.nn.conv import MessagePassing
+
+from .inits import reset
+
+
+class GINConv(MessagePassing):
+    r"""The graph isomorphism operator from the `"How Powerful are
+    Graph Neural Networks?" <https://arxiv.org/abs/1810.00826>`_ paper
+
+    .. math::
+        \mathbf{x}^{\prime}_i = h_{\mathbf{\Theta}} \left( (1 + \epsilon) \cdot
+        \mathbf{x}_i + \sum_{j \in \mathcal{N}(i)} \mathbf{x}_j \right)
+
+    or
+
+    .. math::
+        \mathbf{X}^{\prime} = h_{\mathbf{\Theta}} \left( \left( \mathbf{A} +
+        (1 + \epsilon) \cdot \mathbf{I} \right) \cdot \mathbf{X} \right),
+
+    here :math:`h_{\mathbf{\Theta}}` denotes a neural network, *.i.e.* an MLP.
+
+    Args:
+        nn (torch.nn.Module): A neural network :math:`h_{\mathbf{\Theta}}` that
+            maps node features :obj:`x` of shape :obj:`[-1, in_channels]` to
+            shape :obj:`[-1, out_channels]`, *e.g.*, defined by
+            :class:`torch.nn.Sequential`.
+        eps (float, optional): (Initial) :math:`\epsilon`-value.
+            (default: :obj:`0.`)
+        train_eps (bool, optional): If set to :obj:`True`, :math:`\epsilon`
+            will be a trainable parameter. (default: :obj:`False`)
+        **kwargs (optional): Additional arguments of
+            :class:`torch_geometric.nn.conv.MessagePassing`.
+    """
+    def __init__(self, nn: Callable, eps: float = 0., train_eps: bool = False,
+                 **kwargs):
+        kwargs.setdefault('aggr', 'add')
+        super(GINConv, self).__init__(**kwargs)
+        self.nn = nn
+        self.initial_eps = eps
+        if train_eps:
+            self.eps = torch.nn.Parameter(torch.Tensor([eps]))
+        else:
+            self.register_buffer('eps', torch.Tensor([eps]))
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        reset(self.nn)
+        self.eps.data.fill_(self.initial_eps)
+
+    def forward(self, x: Union[Tensor, OptPairTensor], edge_index: Adj, edge_weight: OptTensor = None,
+                size: Size = None) -> Tensor:
+        """"""
+        if isinstance(x, Tensor):
+            x: OptPairTensor = (x, x)
+
+        # propagate_type: (x: OptPairTensor)
+        out = self.propagate(edge_index, x=x, edge_weight = edge_weight, size=size)
+
+        x_r = x[1]
+        if x_r is not None:
+            out += (1 + self.eps) * x_r
+
+        return self.nn(out)
+
+    #def message(self, x_j: Tensor) -> Tensor:
+    #    return x_j
+    def message(self, x_j: Tensor, edge_weight: OptTensor) -> Tensor:
+        if edge_weight is None:
+            return x_j
+        else:
+            return edge_weight.view(-1, 1) * x_j
+
+
+    def message_and_aggregate(self, adj_t: SparseTensor,
+                              x: OptPairTensor) -> Tensor:
+        adj_t = adj_t.set_value(None, layout=None)
+        return matmul(adj_t, x[0], reduce=self.aggr)
+
+    def __repr__(self):
+        return '{}(nn={})'.format(self.__class__.__name__, self.nn)
+
+
+class GINEConv(MessagePassing):
+    r"""The modified :class:`GINConv` operator from the `"Strategies for
+    Pre-training Graph Neural Networks" <https://arxiv.org/abs/1905.12265>`_
+    paper
+
+    .. math::
+        \mathbf{x}^{\prime}_i = h_{\mathbf{\Theta}} \left( (1 + \epsilon) \cdot
+        \mathbf{x}_i + \sum_{j \in \mathcal{N}(i)} \mathrm{ReLU}
+        ( \mathbf{x}_j + \mathbf{e}_{j,i} ) \right)
+
+    that is able to incorporate edge features :math:`\mathbf{e}_{j,i}` into
+    the aggregation procedure.
+
+    Args:
+        nn (torch.nn.Module): A neural network :math:`h_{\mathbf{\Theta}}` that
+            maps node features :obj:`x` of shape :obj:`[-1, in_channels]` to
+            shape :obj:`[-1, out_channels]`, *e.g.*, defined by
+            :class:`torch.nn.Sequential`.
+        eps (float, optional): (Initial) :math:`\epsilon`-value.
+            (default: :obj:`0.`)
+        train_eps (bool, optional): If set to :obj:`True`, :math:`\epsilon`
+            will be a trainable parameter. (default: :obj:`False`)
+        **kwargs (optional): Additional arguments of
+            :class:`torch_geometric.nn.conv.MessagePassing`.
+    """
+    def __init__(self, nn: Callable, eps: float = 0., train_eps: bool = False,
+                 **kwargs):
+        kwargs.setdefault('aggr', 'add')
+        super(GINEConv, self).__init__(**kwargs)
+        self.nn = nn
+        self.initial_eps = eps
+        if train_eps:
+            self.eps = torch.nn.Parameter(torch.Tensor([eps]))
+        else:
+            self.register_buffer('eps', torch.Tensor([eps]))
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        reset(self.nn)
+        self.eps.data.fill_(self.initial_eps)
+
+    def forward(self, x: Union[Tensor, OptPairTensor], edge_index: Adj,
+                edge_attr: OptTensor = None, size: Size = None) -> Tensor:
+        """"""
+        if isinstance(x, Tensor):
+            x: OptPairTensor = (x, x)
+
+        # Node and edge feature dimensionalites need to match.
+        if isinstance(edge_index, Tensor):
+            assert edge_attr is not None
+            assert x[0].size(-1) == edge_attr.size(-1)
+        elif isinstance(edge_index, SparseTensor):
+            assert x[0].size(-1) == edge_index.size(-1)
+
+        # propagate_type: (x: OptPairTensor, edge_attr: OptTensor)
+        out = self.propagate(edge_index, x=x, edge_attr=edge_attr, size=size)
+
+        x_r = x[1]
+        if x_r is not None:
+            out += (1 + self.eps) * x_r
+
+        return self.nn(out)
+
+    def message(self, x_j: Tensor, edge_attr: Tensor) -> Tensor:
+        return F.relu(x_j + edge_attr)
+
+    def __repr__(self):
+        return '{}(nn={})'.format(self.__class__.__name__, self.nn)

autogl/module/nas/space/gasso_space/inits.py

@@ -0,0 +1,56 @@
+import math
+
+import torch
+
+
+def uniform(size, tensor):
+    if tensor is not None:
+        bound = 1.0 / math.sqrt(size)
+        tensor.data.uniform_(-bound, bound)
+
+
+def kaiming_uniform(tensor, fan, a):
+    if tensor is not None:
+        bound = math.sqrt(6 / ((1 + a**2) * fan))
+        tensor.data.uniform_(-bound, bound)
+
+
+def glorot(tensor):
+    if tensor is not None:
+        stdv = math.sqrt(6.0 / (tensor.size(-2) + tensor.size(-1)))
+        tensor.data.uniform_(-stdv, stdv)
+
+
+def glorot_orthogonal(tensor, scale):
+    if tensor is not None:
+        torch.nn.init.orthogonal_(tensor.data)
+        scale /= ((tensor.size(-2) + tensor.size(-1)) * tensor.var())
+        tensor.data *= scale.sqrt()
+
+
+def zeros(tensor):
+    if tensor is not None:
+        tensor.data.fill_(0)
+
+
+def ones(tensor):
+    if tensor is not None:
+        tensor.data.fill_(1)
+
+
+def normal(tensor, mean, std):
+    if tensor is not None:
+        tensor.data.normal_(mean, std)
+
+
+def reset(nn):
+    def _reset(item):
+        if hasattr(item, 'reset_parameters'):
+            item.reset_parameters()
+
+    if nn is not None:
+        if hasattr(nn, 'children') and len(list(nn.children())) > 0:
+            for item in nn.children():
+                _reset(item)
+        else:
+            _reset(nn)

autogl/module/nas/space/gasso_space/message_passing.jinja

@@ -0,0 +1,153 @@
+from typing import *
+from torch_geometric.typing import *
+
+import torch
+from torch import Tensor
+import torch_sparse
+from torch_sparse import SparseTensor
+from torch_geometric.nn.conv.message_passing import *
+from {{module}} import *
+
+
+class Propagate_{{uid}}(NamedTuple):
+{%- for k, v in prop_types.items() %}
+    {{k}}: {{v}}
+{%- endfor %}
+
+
+class Collect_{{uid}}(NamedTuple):
+{%- for k, v in collect_types.items() %}
+    {{k}}: {{v}}
+{%- endfor %}
+
+
+
+class {{cls_name}}({{parent_cls_name}}):
+
+    @torch.jit._overload_method
+    def __check_input__(self, edge_index, size):
+        # type: (Tensor, Size) -> List[Optional[int]]
+        pass
+
+    @torch.jit._overload_method
+    def __check_input__(self, edge_index, size):
+        # type: (SparseTensor, Size) -> List[Optional[int]]
+        pass
+
+{{check_input}}
+
+    @torch.jit._overload_method
+    def __lift__(self, src, edge_index, dim):
+        # type: (Tensor, Tensor, int) -> Tensor
+        pass
+
+    @torch.jit._overload_method
+    def __lift__(self, src, edge_index, dim):
+        # type: (Tensor, SparseTensor, int) -> Tensor
+        pass
+
+{{lift}}
+
+    @torch.jit._overload_method
+    def __collect__(self, edge_def, size, kwargs):
+        # type: (Tensor, List[Optional[int]], Propagate_{{uid}}) -> Collect_{{uid}}
+        pass
+
+    @torch.jit._overload_method
+    def __collect__(self, edge_def, size, kwargs):
+        # type: (SparseTensor, List[Optional[int]], Propagate_{{uid}}) -> Collect_{{uid}}
+        pass
+
+    def __collect__(self, edge_def, size, kwargs):
+        init = torch.tensor(0.)
+        i, j = (1, 0) if self.flow == 'source_to_target' else (0, 1)
+{% for arg in user_args %}
+{%- if arg[-2:] not in ['_i', '_j'] %}
+        {{arg}} = kwargs.{{arg}}
+{%- else %}
+        {{arg}}: {{collect_types[arg]}} = {% if collect_types[arg][:8] == 'Optional' %}None{% else %}init{% endif %}
+        data = kwargs.{{arg[:-2]}}
+        if isinstance(data, (tuple, list)):
+            assert len(data) == 2
+{%- if arg[-2:] == '_j' %}
+            tmp = data[1]
+            if isinstance(tmp, Tensor):
+                self.__set_size__(size, 1, tmp)
+            {{arg}} = data[0]
+{%- else %}
+            tmp = data[0]
+            if isinstance(tmp, Tensor):
+                self.__set_size__(size, 0, tmp)
+            {{arg}} = data[1]
+{%- endif %}
+        else:
+            {{arg}} = data
+        if isinstance({{arg}}, Tensor):
+            self.__set_size__(size, {% if arg[-2:] == '_j'%}0{% else %}1{% endif %}, {{arg}})
+            {{arg}} = self.__lift__({{arg}}, edge_def, {% if arg[-2:] == "_j" %}j{% else %}i{% endif %})
+{%- endif %}
+{%- endfor %}
+
+        edge_index: Optional[Tensor] = None
+        adj_t: Optional[SparseTensor] = None
+        edge_index_i: torch.Tensor = init
+        edge_index_j: torch.Tensor = init
+        ptr: Optional[Tensor] = None
+        if isinstance(edge_def, Tensor):
+            edge_index = edge_def
+            edge_index_i = edge_def[i]
+            edge_index_j = edge_def[j]
+        elif isinstance(edge_def, SparseTensor):
+            adj_t = edge_def
+            edge_index_i = edge_def.storage.row()
+            edge_index_j = edge_def.storage.col()
+            ptr = edge_def.storage.rowptr()
+            {% if 'edge_weight' in collect_types.keys() %}edge_weight = edge_def.storage.value(){% endif %}
+            {% if 'edge_attr' in collect_types.keys() %}edge_attr = edge_def.storage.value(){% endif %}
+            {% if 'edge_type' in collect_types.keys() %}edge_type = edge_def.storage.value(){% endif %}
+
+        {% if collect_types.get('edge_weight', 'Optional')[:8] != 'Optional' %}assert edge_weight is not None{% endif %}
+        {% if collect_types.get('edge_attr', 'Optional')[:8] != 'Optional' %}assert edge_attr is not None{% endif %}
+        {% if collect_types.get('edge_type', 'Optional')[:8] != 'Optional' %}assert edge_type is not None{% endif %}
+
+        index = edge_index_i
+        size_i = size[1] if size[1] is not None else size[0]
+        size_j = size[0] if size[0] is not None else size[1]
+        dim_size = size_i
+
+        return Collect_{{uid}}({% for k in collect_types.keys() %}{{k}}={{k}}{{ ", " if not loop.last }}{% endfor %})
+
+    @torch.jit._overload_method
+    def propagate(self, edge_index, {{ prop_types.keys()|join(', ') }}, size=None):
+        # type: (Tensor, {{ prop_types.values()|join(', ') }}, Size) -> Tensor
+        pass
+
+    @torch.jit._overload_method
+    def propagate(self, edge_index, {{ prop_types.keys()|join(', ') }}, size=None):
+        # type: (SparseTensor, {{ prop_types.values()|join(', ') }}, Size) -> Tensor
+        pass
+
+    def propagate(self, edge_index, {{ prop_types.keys()|join(', ') }}, size=None):
+        the_size = self.__check_input__(edge_index, size)
+        in_kwargs = Propagate_{{uid}}({% for k in prop_types.keys() %}{{k}}={{k}}{{ ", " if not loop.last }}{% endfor %})
+
+        if self.fuse:
+            if isinstance(edge_index, SparseTensor):
+                out = self.message_and_aggregate(edge_index{% for k in msg_and_aggr_args %}, {{k}}=in_kwargs.{{k}}{% endfor %})
+                return self.update(out{% for k in update_args %}, {{k}}=in_kwargs.{{k}}{% endfor %})
+
+        kwargs = self.__collect__(edge_index, the_size, in_kwargs)
+        out = self.message({% for k in msg_args %}{{k}}=kwargs.{{k}}{{ ", " if not loop.last }}{% endfor %})
+        out = self.aggregate(out{% for k in aggr_args %}, {{k}}=kwargs.{{k}}{% endfor %})
+        return self.update(out{% for k in update_args %}, {{k}}=kwargs.{{k}}{% endfor %})
+
+{%- for (arg_types, return_type_repr) in forward_types %}
+
+    @torch.jit._overload_method
+    {{forward_header}}
+        # type: ({{arg_types|join(', ')}}) -> {{return_type_repr}}
+        pass
+{%- endfor %}
+
+    {{forward_header}}
+{{forward_body}}

autogl/module/nas/space/gasso_space/message_passing.py

@@ -0,0 +1,389 @@
+import os
+import re
+import inspect
+import os.path as osp
+from uuid import uuid1
+from itertools import chain
+from inspect import Parameter
+from typing import List, Optional, Set
+from torch_geometric.typing import Adj, Size
+
+import torch
+from torch import Tensor
+from jinja2 import Template
+from torch_sparse import SparseTensor
+from torch_scatter import gather_csr, scatter, segment_csr
+
+from .utils.helpers import expand_left
+from .utils.jit import class_from_module_repr
+from .utils.typing import (sanitize, split_types_repr, parse_types,
+                           resolve_types)
+from .utils.inspector import Inspector, func_header_repr, func_body_repr
+
+
+class MessagePassing(torch.nn.Module):
+    r"""Base class for creating message passing layers of the form
+
+    .. math::
+        \mathbf{x}_i^{\prime} = \gamma_{\mathbf{\Theta}} \left( \mathbf{x}_i,
+        \square_{j \in \mathcal{N}(i)} \, \phi_{\mathbf{\Theta}}
+        \left(\mathbf{x}_i, \mathbf{x}_j,\mathbf{e}_{j,i}\right) \right),
+
+    where :math:`\square` denotes a differentiable, permutation invariant
+    function, *e.g.*, sum, mean or max, and :math:`\gamma_{\mathbf{\Theta}}`
+    and :math:`\phi_{\mathbf{\Theta}}` denote differentiable functions such as
+    MLPs.
+    See `here <https://pytorch-geometric.readthedocs.io/en/latest/notes/
+    create_gnn.html>`__ for the accompanying tutorial.
+
+    Args:
+        aggr (string, optional): The aggregation scheme to use
+            (:obj:`"add"`, :obj:`"mean"`, :obj:`"max"` or :obj:`None`).
+            (default: :obj:`"add"`)
+        flow (string, optional): The flow direction of message passing
+            (:obj:`"source_to_target"` or :obj:`"target_to_source"`).
+            (default: :obj:`"source_to_target"`)
+        node_dim (int, optional): The axis along which to propagate.
+            (default: :obj:`-2`)
+    """
+
+    special_args: Set[str] = {
+        'edge_index', 'adj_t', 'edge_index_i', 'edge_index_j', 'size',
+        'size_i', 'size_j', 'ptr', 'index', 'dim_size'
+    }
+
+    def __init__(self, aggr: Optional[str] = "add",
+                 flow: str = "source_to_target", node_dim: int = -2):
+
+        super(MessagePassing, self).__init__()
+
+        self.aggr = aggr
+        assert self.aggr in ['add', 'mean', 'max', None]
+
+        self.flow = flow
+        assert self.flow in ['source_to_target', 'target_to_source']
+
+        self.node_dim = node_dim
+
+        self.inspector = Inspector(self)
+        self.inspector.inspect(self.message)
+        self.inspector.inspect(self.aggregate, pop_first=True)
+        self.inspector.inspect(self.message_and_aggregate, pop_first=True)
+        self.inspector.inspect(self.update, pop_first=True)
+
+        self.__user_args__ = self.inspector.keys(
+            ['message', 'aggregate', 'update']).difference(self.special_args)
+        self.__fused_user_args__ = self.inspector.keys(
+            ['message_and_aggregate', 'update']).difference(self.special_args)
+
+        # Support for "fused" message passing.
+        self.fuse = self.inspector.implements('message_and_aggregate')
+
+        # Support for GNNExplainer.
+        self.__explain__ = False
+        self.__edge_mask__ = None
+
+    def __check_input__(self, edge_index, size):
+        the_size: List[Optional[int]] = [None, None]
+
+        if isinstance(edge_index, Tensor):
+            assert edge_index.dtype == torch.long
+            assert edge_index.dim() == 2
+            assert edge_index.size(0) == 2
+            if size is not None:
+                the_size[0] = size[0]
+                the_size[1] = size[1]
+            return the_size
+
+        elif isinstance(edge_index, SparseTensor):
+            if self.flow == 'target_to_source':
+                raise ValueError(
+                    ('Flow direction "target_to_source" is invalid for '
+                     'message propagation via `torch_sparse.SparseTensor`. If '
+                     'you really want to make use of a reverse message '
+                     'passing flow, pass in the transposed sparse tensor to '
+                     'the message passing module, e.g., `adj_t.t()`.'))
+            the_size[0] = edge_index.sparse_size(1)
+            the_size[1] = edge_index.sparse_size(0)
+            return the_size
+
+        raise ValueError(
+            ('`MessagePassing.propagate` only supports `torch.LongTensor` of '
+             'shape `[2, num_messages]` or `torch_sparse.SparseTensor` for '
+             'argument `edge_index`.'))
+
+    def __set_size__(self, size: List[Optional[int]], dim: int, src: Tensor):
+        the_size = size[dim]
+        if the_size is None:
+            size[dim] = src.size(self.node_dim)
+        elif the_size != src.size(self.node_dim):
+            raise ValueError(
+                (f'Encountered tensor with size {src.size(self.node_dim)} in '
+                 f'dimension {self.node_dim}, but expected size {the_size}.'))
+
+    def __lift__(self, src, edge_index, dim):
+        if isinstance(edge_index, Tensor):
+            index = edge_index[dim]
+            return src.index_select(self.node_dim, index)
+        elif isinstance(edge_index, SparseTensor):
+            if dim == 1:
+                rowptr = edge_index.storage.rowptr()
+                rowptr = expand_left(rowptr, dim=self.node_dim, dims=src.dim())
+                return gather_csr(src, rowptr)
+            elif dim == 0:
+                col = edge_index.storage.col()
+                return src.index_select(self.node_dim, col)
+        raise ValueError
+
+    def __collect__(self, args, edge_index, size, kwargs):
+        i, j = (1, 0) if self.flow == 'source_to_target' else (0, 1)
+
+        out = {}
+        for arg in args:
+            if arg[-2:] not in ['_i', '_j']:
+                out[arg] = kwargs.get(arg, Parameter.empty)
+            else:
+                dim = 0 if arg[-2:] == '_j' else 1
+                data = kwargs.get(arg[:-2], Parameter.empty)
+
+                if isinstance(data, (tuple, list)):
+                    assert len(data) == 2
+                    if isinstance(data[1 - dim], Tensor):
+                        self.__set_size__(size, 1 - dim, data[1 - dim])
+                    data = data[dim]
+
+                if isinstance(data, Tensor):
+                    self.__set_size__(size, dim, data)
+                    data = self.__lift__(data, edge_index,
+                                         j if arg[-2:] == '_j' else i)
+
+                out[arg] = data
+
+        if isinstance(edge_index, Tensor):
+            out['adj_t'] = None
+            out['edge_index'] = edge_index
+            out['edge_index_i'] = edge_index[i]
+            out['edge_index_j'] = edge_index[j]
+            out['ptr'] = None
+        elif isinstance(edge_index, SparseTensor):
+            out['adj_t'] = edge_index
+            out['edge_index'] = None
+            out['edge_index_i'] = edge_index.storage.row()
+            out['edge_index_j'] = edge_index.storage.col()
+            out['ptr'] = edge_index.storage.rowptr()
+            out['edge_weight'] = edge_index.storage.value()
+            out['edge_attr'] = edge_index.storage.value()
+            out['edge_type'] = edge_index.storage.value()
+
+        out['index'] = out['edge_index_i']
+        out['size'] = size
+        out['size_i'] = size[1] or size[0]
+        out['size_j'] = size[0] or size[1]
+        out['dim_size'] = out['size_i']
+
+        return out
+
+    def propagate(self, edge_index: Adj, size: Size = None, **kwargs):
+        r"""The initial call to start propagating messages.
+
+        Args:
+            edge_index (Tensor or SparseTensor): A :obj:`torch.LongTensor` or a
+                :obj:`torch_sparse.SparseTensor` that defines the underlying
+                graph connectivity/message passing flow.
+                :obj:`edge_index` holds the indices of a general (sparse)
+                assignment matrix of shape :obj:`[N, M]`.
+                If :obj:`edge_index` is of type :obj:`torch.LongTensor`, its
+                shape must be defined as :obj:`[2, num_messages]`, where
+                messages from nodes in :obj:`edge_index[0]` are sent to
+                nodes in :obj:`edge_index[1]`
+                (in case :obj:`flow="source_to_target"`).
+                If :obj:`edge_index` is of type
+                :obj:`torch_sparse.SparseTensor`, its sparse indices
+                :obj:`(row, col)` should relate to :obj:`row = edge_index[1]`
+                and :obj:`col = edge_index[0]`.
+                The major difference between both formats is that we need to
+                input the *transposed* sparse adjacency matrix into
+                :func:`propagate`.
+            size (tuple, optional): The size :obj:`(N, M)` of the assignment
+                matrix in case :obj:`edge_index` is a :obj:`LongTensor`.
+                If set to :obj:`None`, the size will be automatically inferred
+                and assumed to be quadratic.
+                This argument is ignored in case :obj:`edge_index` is a
+                :obj:`torch_sparse.SparseTensor`. (default: :obj:`None`)
+            **kwargs: Any additional data which is needed to construct and
+                aggregate messages, and to update node embeddings.
+        """
+        size = self.__check_input__(edge_index, size)
+
+        # Run "fused" message and aggregation (if applicable).
+        if (isinstance(edge_index, SparseTensor) and self.fuse
+                and not self.__explain__):
+            coll_dict = self.__collect__(self.__fused_user_args__, edge_index,
+                                         size, kwargs)
+
+            msg_aggr_kwargs = self.inspector.distribute(
+                'message_and_aggregate', coll_dict)
+            out = self.message_and_aggregate(edge_index, **msg_aggr_kwargs)
+
+            update_kwargs = self.inspector.distribute('update', coll_dict)
+            return self.update(out, **update_kwargs)
+
+        # Otherwise, run both functions in separation.
+        elif isinstance(edge_index, Tensor) or not self.fuse:
+            coll_dict = self.__collect__(self.__user_args__, edge_index, size,
+                                         kwargs)
+
+            msg_kwargs = self.inspector.distribute('message', coll_dict)
+            out = self.message(**msg_kwargs)
+
+            # For `GNNExplainer`, we require a separate message and aggregate
+            # procedure since this allows us to inject the `edge_mask` into the
+            # message passing computation scheme.
+            if self.__explain__:
+                edge_mask = self.__edge_mask__.sigmoid()
+                # Some ops add self-loops to `edge_index`. We need to do the
+                # same for `edge_mask` (but do not train those).
+                if out.size(self.node_dim) != edge_mask.size(0):
+                    loop = edge_mask.new_ones(size[0])
+                    edge_mask = torch.cat([edge_mask, loop], dim=0)
+                assert out.size(self.node_dim) == edge_mask.size(0)
+                out = out * edge_mask.view([-1] + [1] * (out.dim() - 1))
+
+            aggr_kwargs = self.inspector.distribute('aggregate', coll_dict)
+            out = self.aggregate(out, **aggr_kwargs)
+
+            update_kwargs = self.inspector.distribute('update', coll_dict)
+            return self.update(out, **update_kwargs)
+
+    def message(self, x_j: Tensor) -> Tensor:
+        r"""Constructs messages from node :math:`j` to node :math:`i`
+        in analogy to :math:`\phi_{\mathbf{\Theta}}` for each edge in
+        :obj:`edge_index`.
+        This function can take any argument as input which was initially
+        passed to :meth:`propagate`.
+        Furthermore, tensors passed to :meth:`propagate` can be mapped to the
+        respective nodes :math:`i` and :math:`j` by appending :obj:`_i` or
+        :obj:`_j` to the variable name, *.e.g.* :obj:`x_i` and :obj:`x_j`.
+        """
+        return x_j
+
+    def aggregate(self, inputs: Tensor, index: Tensor,
+                  ptr: Optional[Tensor] = None,
+                  dim_size: Optional[int] = None) -> Tensor:
+        r"""Aggregates messages from neighbors as
+        :math:`\square_{j \in \mathcal{N}(i)}`.
+
+        Takes in the output of message computation as first argument and any
+        argument which was initially passed to :meth:`propagate`.
+
+        By default, this function will delegate its call to scatter functions
+        that support "add", "mean" and "max" operations as specified in
+        :meth:`__init__` by the :obj:`aggr` argument.
+        """
+        if ptr is not None:
+            ptr = expand_left(ptr, dim=self.node_dim, dims=inputs.dim())
+            return segment_csr(inputs, ptr, reduce=self.aggr)
+        else:
+            return scatter(inputs, index, dim=self.node_dim, dim_size=dim_size,
+                           reduce=self.aggr)
+
+    def message_and_aggregate(self, adj_t: SparseTensor) -> Tensor:
+        r"""Fuses computations of :func:`message` and :func:`aggregate` into a
+        single function.
+        If applicable, this saves both time and memory since messages do not
+        explicitly need to be materialized.
+        This function will only gets called in case it is implemented and
+        propagation takes place based on a :obj:`torch_sparse.SparseTensor`.
+        """
+        raise NotImplementedError
+
+    def update(self, inputs: Tensor) -> Tensor:
+        r"""Updates node embeddings in analogy to
+        :math:`\gamma_{\mathbf{\Theta}}` for each node
+        :math:`i \in \mathcal{V}`.
+        Takes in the output of aggregation as first argument and any argument
+        which was initially passed to :meth:`propagate`.
+        """
+        return inputs
+
+    @torch.jit.unused
+    def jittable(self, typing: Optional[str] = None):
+        r"""Analyzes the :class:`MessagePassing` instance and produces a new
+        jittable module.
+
+        Args:
+            typing (string, optional): If given, will generate a concrete
+                instance with :meth:`forward` types based on :obj:`typing`,
+                *e.g.*: :obj:`"(Tensor, Optional[Tensor]) -> Tensor"`.
+        """
+        # Find and parse `propagate()` types to format `{arg1: type1, ...}`.
+        if hasattr(self, 'propagate_type'):
+            prop_types = {
+                k: sanitize(str(v))
+                for k, v in self.propagate_type.items()
+            }
+        else:
+            source = inspect.getsource(self.__class__)
+            match = re.search(r'#\s*propagate_type:\s*\((.*)\)', source)
+            if match is None:
+                raise TypeError(
+                    'TorchScript support requires the definition of the types '
+                    'passed to `propagate()`. Please specificy them via\n\n'
+                    'propagate_type = {"arg1": type1, "arg2": type2, ... }\n\n'
+                    'or via\n\n'
+                    '# propagate_type: (arg1: type1, arg2: type2, ...)\n\n'
+                    'inside the `MessagePassing` module.')
+            prop_types = split_types_repr(match.group(1))
+            prop_types = dict([re.split(r'\s*:\s*', t) for t in prop_types])
+
+        # Parse `__collect__()` types to format `{arg:1, type1, ...}`.
+        collect_types = self.inspector.types(
+            ['message', 'aggregate', 'update'])
+
+        # Collect `forward()` header, body and @overload types.
+        forward_types = parse_types(self.forward)
+        forward_types = [resolve_types(*types) for types in forward_types]
+        forward_types = list(chain.from_iterable(forward_types))
+
+        keep_annotation = len(forward_types) < 2
+        forward_header = func_header_repr(self.forward, keep_annotation)
+        forward_body = func_body_repr(self.forward, keep_annotation)
+
+        if keep_annotation:
+            forward_types = []
+        elif typing is not None:
+            forward_types = []
+            forward_body = 8 * ' ' + f'# type: {typing}\n{forward_body}'
+
+        root = os.path.dirname(osp.realpath(__file__))
+        with open(osp.join(root, 'message_passing.jinja'), 'r') as f:
+            template = Template(f.read())
+
+        uid = uuid1().hex[:6]
+        cls_name = f'{self.__class__.__name__}Jittable_{uid}'
+        jit_module_repr = template.render(
+            uid=uid,
+            module=str(self.__class__.__module__),
+            cls_name=cls_name,
+            parent_cls_name=self.__class__.__name__,
+            prop_types=prop_types,
+            collect_types=collect_types,
+            user_args=self.__user_args__,
+            forward_header=forward_header,
+            forward_types=forward_types,
+            forward_body=forward_body,
+            msg_args=self.inspector.keys(['message']),
+            aggr_args=self.inspector.keys(['aggregate']),
+            msg_and_aggr_args=self.inspector.keys(['message_and_aggregate']),
+            update_args=self.inspector.keys(['update']),
+            check_input=inspect.getsource(self.__check_input__)[:-1],
+            lift=inspect.getsource(self.__lift__)[:-1],
+        )
+
+        # Instantiate a class from the rendered JIT module representation.
+        cls = class_from_module_repr(cls_name, jit_module_repr)
+        module = cls.__new__(cls)
+        module.__dict__ = self.__dict__.copy()
+        module.jittable = None
+
+        return module

autogl/module/nas/space/gasso_space/sage_conv.py

@@ -0,0 +1,92 @@
+from typing import Union, Tuple
+from torch_geometric.typing import (OptPairTensor, Adj, Size, NoneType,
+                                    OptTensor)
+
+from torch import Tensor
+from torch.nn import Linear
+import torch.nn.functional as F
+from torch_sparse import SparseTensor, matmul
+from torch_geometric.nn.conv import MessagePassing
+
+
+class SAGEConv(MessagePassing):
+    r"""The GraphSAGE operator from the `"Inductive Representation Learning on
+    Large Graphs" <https://arxiv.org/abs/1706.02216>`_ paper
+
+    .. math::
+        \mathbf{x}^{\prime}_i = \mathbf{W}_1 \mathbf{x}_i + \mathbf{W_2} \cdot
+        \mathrm{mean}_{j \in \mathcal{N(i)}} \mathbf{x}_j
+
+    Args:
+        in_channels (int or tuple): Size of each input sample. A tuple
+            corresponds to the sizes of source and target dimensionalities.
+        out_channels (int): Size of each output sample.
+        normalize (bool, optional): If set to :obj:`True`, output features
+            will be :math:`\ell_2`-normalized, *i.e.*,
+            :math:`\frac{\mathbf{x}^{\prime}_i}
+            {\| \mathbf{x}^{\prime}_i \|_2}`.
+            (default: :obj:`False`)
+        bias (bool, optional): If set to :obj:`False`, the layer will not learn
+            an additive bias. (default: :obj:`True`)
+        **kwargs (optional): Additional arguments of
+            :class:`torch_geometric.nn.conv.MessagePassing`.
+    """
+    def __init__(self, in_channels: Union[int, Tuple[int, int]],
+                 out_channels: int, normalize: bool = False,
+                 bias: bool = True, **kwargs):  # yapf: disable
+        kwargs.setdefault('aggr', 'mean')
+        super(SAGEConv, self).__init__(**kwargs)
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.normalize = normalize
+
+        if isinstance(in_channels, int):
+            in_channels = (in_channels, in_channels)
+
+        self.lin_l = Linear(in_channels[0], out_channels, bias=bias)
+        self.lin_r = Linear(in_channels[1], out_channels, bias=False)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        self.lin_l.reset_parameters()
+        self.lin_r.reset_parameters()
+
+    def forward(self, x: Union[Tensor, OptPairTensor], edge_index: Adj, edge_weight: OptTensor = None,
+                size: Size = None) -> Tensor:
+        """"""
+        if isinstance(x, Tensor):
+            x: OptPairTensor = (x, x)
+
+        # propagate_type: (x: OptPairTensor)
+        out = self.propagate(edge_index, x=x, edge_weight = edge_weight, size=size)
+        out = self.lin_l(out)
+
+        x_r = x[1]
+        if x_r is not None:
+            out += self.lin_r(x_r)
+
+        if self.normalize:
+            out = F.normalize(out, p=2., dim=-1)
+
+        return out
+
+    #def message(self, x_j: Tensor) -> Tensor:
+    #    return x_j
+    def message(self, x_j: Tensor, edge_weight: OptTensor) -> Tensor:
+        if edge_weight is None:
+            return x_j
+        else:
+            return edge_weight.view(-1, 1) * x_j
+
+
+
+    def message_and_aggregate(self, adj_t: SparseTensor,
+                              x: OptPairTensor) -> Tensor:
+        adj_t = adj_t.set_value(None, layout=None)
+        return matmul(adj_t, x[0], reduce=self.aggr)
+
+    def __repr__(self):
+        return '{}({}, {})'.format(self.__class__.__name__, self.in_channels,
+                                   self.out_channels)

autogl/module/nas/space/gasso_space/utils/helpers.py

@@ -0,0 +1,7 @@
+import torch
+
+
+def expand_left(src: torch.Tensor, dim: int, dims: int) -> torch.Tensor:
+    for _ in range(dims + dim if dim < 0 else dim):
+        src = src.unsqueeze(0)
+    return src

autogl/module/nas/space/gasso_space/utils/inspector.py

@@ -0,0 +1,86 @@
+import re
+import inspect
+from collections import OrderedDict
+from typing import Dict, List, Any, Optional, Callable, Set
+
+from .typing import parse_types
+
+
+class Inspector(object):
+    def __init__(self, base_class: Any):
+        self.base_class: Any = base_class
+        self.params: Dict[str, Dict[str, Any]] = {}
+
+    def inspect(self, func: Callable,
+                pop_first: bool = False) -> Dict[str, Any]:
+        params = inspect.signature(func).parameters
+        params = OrderedDict(params)
+        if pop_first:
+            params.popitem(last=False)
+        self.params[func.__name__] = params
+
+    def keys(self, func_names: Optional[List[str]] = None) -> Set[str]:
+        keys = []
+        for func in func_names or list(self.params.keys()):
+            keys += self.params[func].keys()
+        return set(keys)
+
+    def __implements__(self, cls, func_name: str) -> bool:
+        if cls.__name__ == 'MessagePassing':
+            return False
+        if func_name in cls.__dict__.keys():
+            return True
+        return any(self.__implements__(c, func_name) for c in cls.__bases__)
+
+    def implements(self, func_name: str) -> bool:
+        return self.__implements__(self.base_class.__class__, func_name)
+
+    def types(self, func_names: Optional[List[str]] = None) -> Dict[str, str]:
+        out: Dict[str, str] = {}
+        for func_name in func_names or list(self.params.keys()):
+            func = getattr(self.base_class, func_name)
+            arg_types = parse_types(func)[0][0]
+            for key in self.params[func_name].keys():
+                if key in out and out[key] != arg_types[key]:
+                    raise ValueError(
+                        (f'Found inconsistent types for argument {key}. '
+                         f'Expected type {out[key]} but found type '
+                         f'{arg_types[key]}.'))
+                out[key] = arg_types[key]
+        return out
+
+    def distribute(self, func_name, kwargs: Dict[str, Any]):
+        out = {}
+        for key, param in self.params[func_name].items():
+            data = kwargs.get(key, inspect.Parameter.empty)
+            if data is inspect.Parameter.empty:
+                if param.default is inspect.Parameter.empty:
+                    raise TypeError(f'Required parameter {key} is empty.')
+                data = param.default
+            out[key] = data
+        return out
+
+
+def func_header_repr(func: Callable, keep_annotation: bool = True) -> str:
+    source = inspect.getsource(func)
+    signature = inspect.signature(func)
+
+    if keep_annotation:
+        return ''.join(re.split(r'(\).*?:.*?\n)', source,
+                                maxsplit=1)[:2]).strip()
+
+    params_repr = ['self']
+    for param in signature.parameters.values():
+        params_repr.append(param.name)
+        if param.default is not inspect.Parameter.empty:
+            params_repr[-1] += f'={param.default}'
+
+    return f'def {func.__name__}({", ".join(params_repr)}):'
+
+
+def func_body_repr(func: Callable, keep_annotation: bool = True) -> str:
+    source = inspect.getsource(func)
+    body_repr = re.split(r'\).*?:.*?\n', source, maxsplit=1)[1]
+    if not keep_annotation:
+        body_repr = re.sub(r'\s*# type:.*\n', '', body_repr)
+    return body_repr

autogl/module/nas/space/gasso_space/utils/jit.py

@@ -0,0 +1,19 @@
+import sys
+import os.path as osp
+from getpass import getuser
+from tempfile import NamedTemporaryFile as TempFile, gettempdir
+from importlib.util import module_from_spec, spec_from_file_location
+
+from torch_geometric.data.makedirs import makedirs
+
+
+def class_from_module_repr(cls_name, module_repr):
+    path = osp.join(gettempdir(), f'{getuser()}_pyg_jit')
+    makedirs(path)
+    with TempFile(mode='w+', suffix='.py', delete=False, dir=path) as f:
+        f.write(module_repr)
+    spec = spec_from_file_location(cls_name, f.name)
+    mod = module_from_spec(spec)
+    sys.modules[cls_name] = mod
+    spec.loader.exec_module(mod)
+    return getattr(mod, cls_name)

autogl/module/nas/space/gasso_space/utils/typing.py

@@ -0,0 +1,107 @@
+import re
+import inspect
+import pyparsing as pp
+from itertools import product
+from collections import OrderedDict
+from typing import Callable, Tuple, Dict, List
+
+
+def split_types_repr(types_repr: str) -> List[str]:
+    out = []
+    i = depth = 0
+    for j, char in enumerate(types_repr):
+        if char == '[':
+            depth += 1
+        elif char == ']':
+            depth -= 1
+        elif char == ',' and depth == 0:
+            out.append(types_repr[i:j].strip())
+            i = j + 1
+    out.append(types_repr[i:].strip())
+    return out
+
+
+def sanitize(type_repr: str):
+    type_repr = re.sub(r'<class \'(.*)\'>', r'\1', type_repr)
+    type_repr = type_repr.replace('typing.', '')
+    type_repr = type_repr.replace('torch_sparse.tensor.', '')
+    type_repr = type_repr.replace('Adj', 'Union[Tensor, SparseTensor]')
+
+    # Replace `Union[..., NoneType]` by `Optional[...]`.
+    sexp = pp.nestedExpr(opener='[', closer=']')
+    tree = sexp.parseString(f'[{type_repr.replace(",", " ")}]').asList()[0]
+
+    def union_to_optional_(tree):
+        for i in range(len(tree)):
+            e, n = tree[i], tree[i + 1] if i + 1 < len(tree) else []
+            if e == 'Union' and n[-1] == 'NoneType':
+                tree[i] = 'Optional'
+                tree[i + 1] = tree[i + 1][:-1]
+            elif e == 'Union' and 'NoneType' in n:
+                idx = n.index('NoneType')
+                n[idx] = [n[idx - 1]]
+                n[idx - 1] = 'Optional'
+            elif isinstance(e, list):
+                tree[i] = union_to_optional_(e)
+        return tree
+
+    tree = union_to_optional_(tree)
+    type_repr = re.sub(r'\'|\"', '', str(tree)[1:-1]).replace(', [', '[')
+
+    return type_repr
+
+
+def param_type_repr(param) -> str:
+    if param.annotation is inspect.Parameter.empty:
+        return 'torch.Tensor'
+    return sanitize(re.split(r':|='.strip(), str(param))[1])
+
+
+def return_type_repr(signature) -> str:
+    return_type = signature.return_annotation
+    if return_type is inspect.Parameter.empty:
+        return 'torch.Tensor'
+    elif str(return_type)[:6] != '<class':
+        return sanitize(str(return_type))
+    elif return_type.__module__ == 'builtins':
+        return return_type.__name__
+    else:
+        return f'{return_type.__module__}.{return_type.__name__}'
+
+
+def parse_types(func: Callable) -> List[Tuple[Dict[str, str], str]]:
+    source = inspect.getsource(func)
+    signature = inspect.signature(func)
+
+    # Parse `# type: (...) -> ...` annotation. Note that it is allowed to pass
+    # multiple `# type:` annotations in `forward()`.
+    iterator = re.finditer(r'#\s*type:\s*\((.*)\)\s*->\s*(.*)\s*\n', source)
+    matches = list(iterator)
+
+    if len(matches) > 0:
+        out = []
+        args = list(signature.parameters.keys())
+        for match in matches:
+            arg_types_repr, return_type = match.groups()
+            arg_types = split_types_repr(arg_types_repr)
+            arg_types = OrderedDict((k, v) for k, v in zip(args, arg_types))
+            return_type = return_type.split('#')[0].strip()
+            out.append((arg_types, return_type))
+        return out
+
+    # Alternatively, parse annotations using the inspected signature.
+    else:
+        ps = signature.parameters
+        arg_types = OrderedDict((k, param_type_repr(v)) for k, v in ps.items())
+        return [(arg_types, return_type_repr(signature))]
+
+
+def resolve_types(arg_types: Dict[str, str],
+                  return_type_repr: str) -> List[Tuple[List[str], str]]:
+    out = []
+    for type_repr in arg_types.values():
+        if type_repr[:5] == 'Union':
+            out.append(split_types_repr(type_repr[6:-1]))
+        else:
+            out.append([type_repr])
+    return [(x, return_type_repr) for x in product(*out)]

autogl/module/nas/space/graph_nas.py

@@ -208,4 +208,4 @@ class GraphNasNodeClassificationSpace(BaseSpace):
 
     def parse_model(self, selection, device) -> BaseAutoModel:
         # return AutoGCN(self.input_dim, self.output_dim, device)
-        return self.wrap(device).fix(selection)
+        return self.wrap().fix(selection)

autogl/module/nas/space/graph_nas_macro.py

@@ -10,6 +10,307 @@ from .operation import act_map
 from ..utils import count_parameters, measure_latency
 
 from ..backend import *
+# import dgl
+# from dgl import function as fn
+
+
+special_args = [
+    "edge_index",
+    "edge_index_i",
+    "edge_index_j",
+    "size",
+    "size_i",
+    "size_j",
+]
+__size_error_msg__ = (
+    "All tensors which should get mapped to the same source "
+    "or target nodes must be of same size in dimension 0."
+)
+
+is_python2 = sys.version_info[0] < 3
+getargspec = inspect.getargspec if is_python2 else inspect.getfullargspec
+
+
+def scatter_(name, src, index, dim_size=None):
+    r"""Aggregates all values from the :attr:`src` tensor at the indices
+    specified in the :attr:`index` tensor along the first dimension.
+    If multiple indices reference the same location, their contributions
+    are aggregated according to :attr:`name` (either :obj:`"add"`,
+    :obj:`"mean"` or :obj:`"max"`).
+
+    Args:
+        name (string): The aggregation to use (:obj:`"add"`, :obj:`"mean"`,
+            :obj:`"max"`).
+        src (Tensor): The source tensor.
+        index (LongTensor): The indices of elements to scatter.
+        dim_size (int, optional): Automatically create output tensor with size
+            :attr:`dim_size` in the first dimension. If set to :attr:`None`, a
+            minimal sized output tensor is returned. (default: :obj:`None`)
+
+    :rtype: :class:`Tensor`
+    """
+
+    assert name in ["add", "mean", "max"]
+
+    op = getattr(torch_scatter, "scatter_{}".format(name))
+    fill_value = -1e9 if name == "max" else 0
+
+    out = op(src, index, 0, None, dim_size)
+    if isinstance(out, tuple):
+        out = out[0]
+
+    if name == "max":
+        out[out == fill_value] = 0
+
+    return out
+
+
+class MessagePassing(torch.nn.Module):
+    def __init__(self, aggr="add", flow="source_to_target"):
+        super(MessagePassing, self).__init__()
+
+        self.aggr = aggr
+        assert self.aggr in ["add", "mean", "max"]
+
+        self.flow = flow
+        assert self.flow in ["source_to_target", "target_to_source"]
+
+        self.__message_args__ = getargspec(self.message)[0][1:]
+        self.__special_args__ = [
+            (i, arg)
+            for i, arg in enumerate(self.__message_args__)
+            if arg in special_args
+        ]
+        self.__message_args__ = [
+            arg for arg in self.__message_args__ if arg not in special_args
+        ]
+        self.__update_args__ = getargspec(self.update)[0][2:]
+
+    def propagate(self, edge_index, size=None, **kwargs):
+        r"""The initial call to start propagating messages.
+
+        Args:
+            edge_index (Tensor): The indices of a general (sparse) assignment
+                matrix with shape :obj:`[N, M]` (can be directed or
+                undirected).
+            size (list or tuple, optional): The size :obj:`[N, M]` of the
+                assignment matrix. If set to :obj:`None`, the size is tried to
+                get automatically inferrred. (default: :obj:`None`)
+            **kwargs: Any additional data which is needed to construct messages
+                and to update node embeddings.
+        """
+
+        size = [None, None] if size is None else list(size)
+        assert len(size) == 2
+
+        i, j = (0, 1) if self.flow == "target_to_source" else (1, 0)
+        ij = {"_i": i, "_j": j}
+
+        message_args = []
+        for arg in self.__message_args__:
+            if arg[-2:] in ij.keys():
+                tmp = kwargs.get(arg[:-2], None)
+                if tmp is None:  # pragma: no cover
+                    message_args.append(tmp)
+                else:
+                    idx = ij[arg[-2:]]
+                    if isinstance(tmp, tuple) or isinstance(tmp, list):
+                        assert len(tmp) == 2
+                        if tmp[1 - idx] is not None:
+                            if size[1 - idx] is None:
+                                size[1 - idx] = tmp[1 - idx].size(0)
+                            if size[1 - idx] != tmp[1 - idx].size(0):
+                                raise ValueError(__size_error_msg__)
+                        tmp = tmp[idx]
+
+                    if size[idx] is None:
+                        size[idx] = tmp.size(0)
+                    if size[idx] != tmp.size(0):
+                        raise ValueError(__size_error_msg__)
+
+                    tmp = torch.index_select(tmp, 0, edge_index[idx])
+                    message_args.append(tmp)
+            else:
+                message_args.append(kwargs.get(arg, None))
+
+        size[0] = size[1] if size[0] is None else size[0]
+        size[1] = size[0] if size[1] is None else size[1]
+
+        kwargs["edge_index"] = edge_index
+        kwargs["size"] = size
+
+        for (idx, arg) in self.__special_args__:
+            if arg[-2:] in ij.keys():
+                message_args.insert(idx, kwargs[arg[:-2]][ij[arg[-2:]]])
+            else:
+                message_args.insert(idx, kwargs[arg])
+
+        update_args = [kwargs[arg] for arg in self.__update_args__]
+
+        out = self.message(*message_args)
+        if self.aggr in ["add", "mean", "max"]:
+            out = scatter_(self.aggr, out, edge_index[i], dim_size=size[i])
+        else:
+            pass
+        out = self.update(out, *update_args)
+
+        return out
+
+    def message(self, x_j):  # pragma: no cover
+        r"""Constructs messages in analogy to :math:`\phi_{\mathbf{\Theta}}`
+        for each edge in :math:`(i,j) \in \mathcal{E}`.
+        Can take any argument which was initially passed to :meth:`propagate`.
+        In addition, features can be lifted to the source node :math:`i` and
+        target node :math:`j` by appending :obj:`_i` or :obj:`_j` to the
+        variable name, *.e.g.* :obj:`x_i` and :obj:`x_j`."""
+
+        return x_j
+
+    def update(self, aggr_out):  # pragma: no cover
+        r"""Updates node embeddings in analogy to
+        :math:`\gamma_{\mathbf{\Theta}}` for each node
+        :math:`i \in \mathcal{V}`.
+        Takes in the output of aggregation as first argument and any argument
+        which was initially passed to :meth:`propagate`."""
+
+        return aggr_out
+
+
+class GeoLayerPYG(MessagePassing):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        heads=1,
+        concat=True,
+        negative_slope=0.2,
+        dropout=0,
+        bias=True,
+        att_type="gat",
+        agg_type="sum",
+        pool_dim=0,
+    ):
+        if agg_type in ["sum", "mlp"]:
+            super(GeoLayerPYG, self).__init__("add")
+        elif agg_type in ["mean", "max"]:
+            super(GeoLayerPYG, self).__init__(agg_type)
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.heads = heads
+        self.concat = concat
+        self.negative_slope = negative_slope
+        self.dropout = dropout
+        self.att_type = att_type
+        self.agg_type = agg_type
+
+        # GCN weight
+        self.gcn_weight = None
+
+        self.weight = Parameter(torch.Tensor(in_channels, heads * out_channels))
+        self.att = Parameter(torch.Tensor(1, heads, 2 * out_channels))
+
+        if bias and concat:
+            self.bias = Parameter(torch.Tensor(heads * out_channels))
+        elif bias and not concat:
+            self.bias = Parameter(torch.Tensor(out_channels))
+        else:
+            self.register_parameter("bias", None)
+
+        if self.att_type in ["generalized_linear"]:
+            self.general_att_layer = torch.nn.Linear(out_channels, 1, bias=False)
+
+        if self.agg_type in ["mean", "max", "mlp"]:
+            if pool_dim <= 0:
+                pool_dim = 128
+        self.pool_dim = pool_dim
+        if pool_dim != 0:
+            self.pool_layer = torch.nn.ModuleList()
+            self.pool_layer.append(torch.nn.Linear(self.out_channels, self.pool_dim))
+            self.pool_layer.append(torch.nn.Linear(self.pool_dim, self.out_channels))
+        else:
+            pass
+        self.reset_parameters()
+
+    @staticmethod
+    def norm(edge_index, num_nodes, edge_weight, improved=False, dtype=None):
+        if edge_weight is None:
+            edge_weight = torch.ones(
+                (edge_index.size(1),), dtype=dtype, device=edge_index.device
+            )
+
+        fill_value = 1 if not improved else 2
+        edge_index, edge_weight = add_remaining_self_loops(
+            edge_index, edge_weight, fill_value, num_nodes
+        )
+
+        row, col = edge_index
+        deg = scatter_add(edge_weight, row, dim=0, dim_size=num_nodes)
+        deg_inv_sqrt = deg.pow(-0.5)
+        deg_inv_sqrt[deg_inv_sqrt == float("inf")] = 0
+
+        return edge_index, deg_inv_sqrt[row] * edge_weight * deg_inv_sqrt[col]
+
+    def reset_parameters(self):
+        glorot(self.weight)
+        glorot(self.att)
+        zeros(self.bias)
+
+        if self.att_type in ["generalized_linear"]:
+            glorot(self.general_att_layer.weight)
+
+        if self.pool_dim != 0:
+            for layer in self.pool_layer:
+                glorot(layer.weight)
+                zeros(layer.bias)
+
+    def forward(self, x, edge_index):
+        """"""
+        edge_index, _ = remove_self_loops(edge_index)
+        edge_index, _ = add_self_loops(edge_index, num_nodes=x.size(0))
+        # prepare
+        x = torch.mm(x, self.weight).view(-1, self.heads, self.out_channels)
+        # x [2708,2,4] weight [1433,8]
+        return self.propagate(edge_index, x=x, num_nodes=x.size(0))
+
+    def message(self, x_i, x_j, edge_index, num_nodes):
+
+        # x_i torch.Size([13264, 2, 4])
+        # x_j torch.Size([13264, 2, 4])
+        # edge_index torch.Size([2, 13264])
+        # num_nodes 2708
+        if self.att_type == "const":
+            if self.training and self.dropout > 0:
+                x_j = F.dropout(x_j, p=self.dropout, training=True)
+            neighbor = x_j
+        elif self.att_type == "gcn":
+            if self.gcn_weight is None or self.gcn_weight.size(0) != x_j.size(
+                0
+            ):  # 对于不同的图gcn_weight需要重新计算
+                _, norm = self.norm(edge_index, num_nodes, None)
+                self.gcn_weight = norm
+            neighbor = self.gcn_weight.view(-1, 1, 1) * x_j
+        else:
+            # Compute attention coefficients.
+            alpha = self.apply_attention(edge_index, num_nodes, x_i, x_j)
+            alpha = softmax(alpha, edge_index[0], num_nodes=num_nodes)
+            # Sample attention coefficients stochastically.
+            if self.training and self.dropout > 0:
+                alpha = F.dropout(alpha, p=self.dropout, training=True)
+
+            neighbor = x_j * alpha.view(-1, self.heads, 1)
+        # pool_layer
+        # (0): Linear(in_features=4, out_features=128, bias=True)
+        #   (1): Linear(in_features=128, out_features=4, bias=True)
+        if self.pool_dim > 0:
+            # neighbor torch.Size([13264, 2, 4])
+            for layer in self.pool_layer:
+                neighbor = layer(neighbor)
+        return neighbor
+
+    def apply_attention(self, edge_index, num_nodes, x_i, x_j):
+        if self.att_type == "gat":
+            alpha = (torch.cat([x_i, x_j], dim=-1) * self.att).sum(dim=-1)
+            alpha = F.leaky_relu(alpha, self.negative_slope)
 
 from operator import *
 from .operation import *
@@ -183,7 +484,7 @@ class GraphNasMacroNodeClassificationSpace(BaseSpace):
             multi_label=False,
             batch_normal=False,
             layers=self.layer_number,
-        ).wrap(device)
+        ).wrap()
         return model
 
 

autogl/module/nas/utils.py

@@ -63,8 +63,7 @@ class PathSamplingLayerChoice(nn.Module):
         return _get_mask(self.sampled, len(self))
 
     def __repr__(self):
-        return f"PathSamplingLayerChoice(chosen={self.sampled},{super().__repr__()})"
-
+        return f"PathSamplingLayerChoice(op_names={self.op_names}, chosen={self.sampled})"
 
 class PathSamplingInputChoice(nn.Module):
     """
@@ -104,6 +103,17 @@ class PathSamplingInputChoice(nn.Module):
 
 
 def get_hardware_aware_metric(model, hardware_metric):
+    """
+    Get architectures' hardware-aware metrics
+
+    Attributes
+    ----------
+    model : BaseSpace
+        The architecture to be evaluated
+    hardware_metric : str
+        The name of hardware-aware metric. Can be 'parameter' or 'latency'
+    """
+
     if hardware_metric == 'parameter':
         return count_parameters(model)
     elif hardware_metric == 'latency':

configs/nodeclf_nas_gasso.yml

@@ -0,0 +1,40 @@
+ensemble:
+  name: null
+feature:
+- name: NormalizeFeatures
+hpo:
+  max_evals: 10
+  name: random
+nas:
+  space:
+    name: gassospace
+    hidden_dim: 64
+    layer_number: 2
+  algorithm:
+    name: gasso
+    num_epochs: 250
+  estimator:
+    name: oneshot
+models: []
+trainer:
+  hp_space:
+  - maxValue: 300
+    minValue: 100
+    parameterName: max_epoch
+    scalingType: LINEAR
+    type: INTEGER
+  - maxValue: 30
+    minValue: 10
+    parameterName: early_stopping_round
+    scalingType: LINEAR
+    type: INTEGER
+  - maxValue: 0.05
+    minValue: 0.01
+    parameterName: lr
+    scalingType: LOG
+    type: DOUBLE
+  - maxValue: 0.0005
+    minValue: 5.0e-05
+    parameterName: weight_decay
+    scalingType: LOG
+    type: DOUBLE

docs/docfile/tutorial/t_nas.rst

@@ -4,7 +4,25 @@ Neural Architecture Search
 ============================
 
 We support different neural architecture search algorithm in variant search space.
-To be more flexible, we modulize NAS process with three part: algorithm, space and estimator.
+Neural architecture search is usually constructed by three modules: search space, search strategy and estimation strategy.
+
+The search space describes all possible architectures to be searched. There are mainly two parts of the space formulated, the operations(e.g. GCNconv, GATconv) and the input-ouput realations.
+A large space may have better optimal architecture but demands more effect to explore.
+Human knowledge can help to design a reasonable search space to reduce the efforts of search strategy.
+
+The search strategy controls how to explore the search sapce. 
+It encompasses the classical exploration-exploitation trade-off since.
+On the one hand, it is desirable to find well-performing architectures quickly, 
+while on the other hand, premature convergence to a region of suboptimal architectures should be avoided.
+
+The estimation strategy gives the performance of certain architectures when it is explored.
+The simplest option is to perform a standard training and validation of the architecture on data.
+Since there are lots of architectures need estimating in the whole searching process, estimation strategy is desired to be very efficient to save computational resources.
+
+.. image:: ../resources/nas.svg
+   :align: center
+
+To be more flexible, we modulize NAS process with three part: algorithm, space and estimator, corresponding to the three module search space, search strategy and estimation strategy.
 Different models in different parts can be composed in some certain constrains.
 If you want to design your own NAS process, you can change any of those parts according to your demand.
 
@@ -99,7 +117,7 @@ Here is an example.
 
         # For one-shot fashion, you can directly use following scheme in ``parse_model``
         def parse_model(self, selection, device) -> BaseModel:
-            return self.wrap(device).fix(selection)
+            return self.wrap().fix(selection)
 
 Also, you can use the way which does not support one shot fashion.
 In this way, you can directly copy you model with few changes.
@@ -135,7 +153,7 @@ But you can only use sample-based search strategy.
         # For non-one-shot fashion, you can directly return your model based on the choices
         # ``YourModel`` must inherit BaseSpace.
         def parse_model(self, selection, device) -> BaseModel:
-            model = YourModel(selection, self.input_dim, self.output_dim).wrap(device)
+            model = YourModel(selection, self.input_dim, self.output_dim).wrap()
             return model
 
     # YourModel can be defined as follows

examples/gasso_test.py

@@ -0,0 +1,25 @@
+import os
+os.environ["AUTOGL_BACKEND"] = "pyg"
+import sys
+sys.path.append('../')
+from autogl.datasets import build_dataset_from_name
+from autogl.solver import AutoNodeClassifier
+from autogl.module.train import Acc
+from autogl.solver.utils import set_seed
+import argparse
+
+if __name__ == '__main__':
+    set_seed(202106)
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--config', type=str, default='../configs/nodeclf_nas_gasso.yml')
+    parser.add_argument('--dataset', choices=['cora', 'citeseer', 'pubmed'], default='citeseer', type=str)
+
+    args = parser.parse_args()
+
+    #dataset = build_dataset_from_name(args.dataset, path = "/DATA/DATANAS1/qinyj/enhgnas/")
+    dataset = build_dataset_from_name(args.dataset, path = "~/AGL/")
+    solver = AutoNodeClassifier.from_config(args.config)
+    solver.fit(dataset)
+    solver.get_leaderboard().show()
+    out = solver.predict_proba()
+    print('acc on dataset', Acc.evaluate(out, dataset[0].y[dataset[0].test_mask].detach().numpy()))

examples/graphnas.py

@@ -12,7 +12,7 @@ if __name__ == '__main__':
 
     args = parser.parse_args()
 
-    dataset = build_dataset_from_name('cora')
+    dataset = build_dataset_from_name(args.dataset)
     label = dataset[0].nodes.data["y" if DependentBackend.is_pyg() else "label"][dataset[0].nodes.data["test_mask"]].cpu().numpy()
     solver = AutoNodeClassifier.from_config(args.config)
     solver.fit(dataset)

test/nas/node_classification.py

@@ -27,6 +27,7 @@ from autogl.module.nas.space.single_path import SinglePathNodeClassificationSpac
 from autogl.module.nas.space.graph_nas import GraphNasNodeClassificationSpace
 from autogl.module.nas.space.graph_nas_macro import GraphNasMacroNodeClassificationSpace
 from autogl.module.nas.estimator.one_shot import OneShotEstimator
+from autogl.module.nas.space.autoattend import AutoAttendNodeClassificationSpace
 from autogl.module.nas.backend import bk_feat, bk_label
 from autogl.module.nas.algorithm import Darts, RL, GraphNasRL, Enas, RandomSearch,Spos
 import numpy as np
@@ -103,39 +104,56 @@ if __name__ == "__main__":
     space = GraphNasNodeClassificationSpace().cuda()
     space.instantiate(input_dim=di, output_dim=do)
     esti = OneShotEstimator()
-    algo = RandomSearch(num_epochs=10)
+    algo = RandomSearch(num_epochs=100)
     model = algo.search(space, dataset, esti)
     test_model(model, data, True)
 
-    print("Random search + singlepath ")
-    space = SinglePathNodeClassificationSpace().cuda()
+    print("Random search + AutoAttend ")
+    space = AutoAttendNodeClassificationSpace().cuda()
     space.instantiate(input_dim=di, output_dim=do)
     esti = OneShotEstimator()
     algo = RandomSearch(num_epochs=10)
     model = algo.search(space, dataset, esti)
+    print(model)
     test_model(model, data, True)
 
-    print("rl + graphnas ")
-    space = GraphNasNodeClassificationSpace().cuda()
+    print("rl + AutoAttend ")
+    space = AutoAttendNodeClassificationSpace().cuda()
     space.instantiate(input_dim=di, output_dim=do)
     esti = OneShotEstimator()
     algo = RL(num_epochs=10)
     model = algo.search(space, dataset, esti)
     test_model(model, data, True)
 
-    print("graphnasrl + graphnas ")
+    print("darts + graphnas ")
+    space = AutoAttendNodeClassificationSpace().cuda()
+    space.instantiate(input_dim=di, output_dim=do)
+    esti = OneShotEstimator()
+    algo = Darts(num_epochs=10)
+    model = algo.search(space, dataset, esti)
+    test_model(model, data, True)
+
+    print("Random search + graphnas ")
     space = GraphNasNodeClassificationSpace().cuda()
     space.instantiate(input_dim=di, output_dim=do)
     esti = OneShotEstimator()
-    algo = GraphNasRL(num_epochs=10)
+    algo = RandomSearch(num_epochs=10)
+    model = algo.search(space, dataset, esti)
+    test_model(model, data, True)
+
+    print("rl + graphnas ")
+    space = GraphNasNodeClassificationSpace().cuda()
+    space.instantiate(input_dim=di, output_dim=do)
+    esti = OneShotEstimator()
+    algo = RL(num_epochs=10)
     model = algo.search(space, dataset, esti)
     test_model(model, data, True)
 
-    print("enas + graphnas ")
+    print("graphnasrl + graphnas ")
     space = GraphNasNodeClassificationSpace().cuda()
     space.instantiate(input_dim=di, output_dim=do)
     esti = OneShotEstimator()
-    algo = Enas(num_epochs=10)
+    algo = GraphNasRL(num_epochs=10)
     model = algo.search(space, dataset, esti)
     test_model(model, data, True)