def __init__(self, keys: List[Union[NodeType, EdgeType]],
out_channels: int):
super().__init__()
self.out_channels = out_channels
self.lins = torch.nn.ModuleDict()
for key in keys:
self.lins[key2str(key)] = Linear(-1, out_channels, bias=False)
def __init__(
self,
in_channels: Union[int, Dict[str, int]],
out_channels: int,
metadata: Metadata,
heads: int = 1,
group: str = "sum",
**kwargs,
):
super().__init__(aggr='add', node_dim=0, **kwargs)
if not isinstance(in_channels, dict):
in_channels = {node_type: in_channels for node_type in metadata[0]}
self.in_channels = in_channels
self.out_channels = out_channels
self.heads = heads
self.group = group
self.k_lin = torch.nn.ModuleDict()
self.q_lin = torch.nn.ModuleDict()
self.v_lin = torch.nn.ModuleDict()
self.a_lin = torch.nn.ModuleDict()
self.skip = torch.nn.ParameterDict()
for node_type, in_channels in self.in_channels.items():
self.k_lin[node_type] = Linear(in_channels, out_channels)
self.q_lin[node_type] = Linear(in_channels, out_channels)
self.v_lin[node_type] = Linear(in_channels, out_channels)
self.a_lin[node_type] = Linear(out_channels, out_channels)
self.skip[node_type] = Parameter(torch.Tensor(1))
self.a_rel = torch.nn.ParameterDict()
self.m_rel = torch.nn.ParameterDict()
self.p_rel = torch.nn.ParameterDict()
dim = out_channels // heads
for edge_type in metadata[1]:
edge_type = '__'.join(edge_type)
self.a_rel[edge_type] = Parameter(torch.Tensor(heads, dim, dim))
self.m_rel[edge_type] = Parameter(torch.Tensor(heads, dim, dim))
self.p_rel[edge_type] = Parameter(torch.Tensor(heads))
self.reset_parameters()
def __init__(
self,
in_channels: Union[int, Dict[str, int]],
out_channels: int,
metadata: Metadata,
heads: int = 1,
negative_slope=0.2,
dropout: float = 0.0,
**kwargs,
):
super().__init__(aggr='add', node_dim=0, **kwargs)
if not isinstance(in_channels, dict):
in_channels = {node_type: in_channels for node_type in metadata[0]}
self.heads = heads
self.in_channels = in_channels
self.out_channels = out_channels
self.negative_slope = negative_slope
self.metadata = metadata
self.dropout = dropout
self.k_lin = nn.Linear(out_channels, out_channels)
self.q = nn.Parameter(torch.Tensor(1, out_channels))
self.proj = nn.ModuleDict()
for node_type, in_channels in self.in_channels.items():
self.proj[node_type] = Linear(in_channels, out_channels)
self.lin_src = nn.ParameterDict()
self.lin_dst = nn.ParameterDict()
dim = out_channels // heads
for edge_type in metadata[1]:
edge_type = '__'.join(edge_type)
self.lin_src[edge_type] = nn.Parameter(torch.Tensor(1, heads, dim))
self.lin_dst[edge_type] = nn.Parameter(torch.Tensor(1, heads, dim))
self.reset_parameters()