k3-node 1.0.0__py3-none-any.whl

k3_node/layers/conv/han_conv.py

@@ -0,0 +1,175 @@
+from typing import Dict, List, Optional, Tuple, Union
+import keras
+from keras import ops
+from k3_node.layers.conv.message_passing import MessagePassing
+from k3_node.layers.conv.utils import softmax
+
+
+def semantic_group(xs: List[any], q: any, k_lin: keras.layers.Layer) -> Tuple[Optional[any], Optional[any]]:
+    if len(xs) == 0:
+        return None, None
+    num_edge_types = len(xs)
+    out = ops.stack(xs, axis=0)  # (E, N, C)
+    if ops.size(out) == 0:
+        return ops.reshape(out, (0, ops.shape(out)[-1])), None
+
+    k_proj = ops.tanh(k_lin(out))  # (E, N, C)
+    mean_k = ops.mean(k_proj, axis=1)  # (E, C)
+    attn_score = ops.sum(q * mean_k, axis=-1)  # (E,)
+    if num_edge_types == 1:
+        attn = ops.ones_like(attn_score)
+    else:
+        attn = ops.softmax(attn_score, axis=0)  # (E,)
+
+    attn_expanded = ops.reshape(attn, (num_edge_types, 1, 1))
+    out = ops.sum(attn_expanded * out, axis=0)  # (N, C)
+    return out, attn
+
+
+class HANConv(MessagePassing):
+    r"""The Heterogeneous Graph Attention Operator from the
+    `"Heterogeneous Graph Attention Network" <https://arxiv.org/abs/1903.07293>`_ paper.
+
+    Args:
+        in_channels (int or Dict[str, int]): Size of each input sample of every node type.
+        out_channels (int): Size of each output sample.
+        metadata (Tuple[List[str], List[Tuple[str, str, str]]]): Node types and edge types.
+        heads (int, optional): Number of multi-head-attentions. (default: :obj:`1`)
+        negative_slope (float, optional): LeakyReLU angle of the negative slope. (default: :obj:`0.2`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import HANConv
+
+        x_dict = {
+            "author": np.random.rand(3, 8).astype("float32"),
+            "paper": np.random.rand(4, 8).astype("float32"),
+        }
+        edge_index_dict = {
+            ("author", "writes", "paper"): np.array([[0, 1, 2], [1, 2, 3]]),
+            ("paper", "written_by", "author"): np.array([[1, 2, 3], [0, 1, 2]]),
+        }
+        metadata = (list(x_dict), list(edge_index_dict))  # (node types, edge types)
+
+        layer = HANConv(in_channels={"author": 8, "paper": 8}, out_channels=16, metadata=metadata, heads=2)
+        out_dict = layer(x_dict, edge_index_dict)
+        print(tuple(out_dict["author"].shape), tuple(out_dict["paper"].shape))  # (3, 16) (4, 16)
+        ```
+    """
+
+    def __init__(
+        self,
+        in_channels: Union[int, Dict[str, int]],
+        out_channels: int,
+        metadata: Tuple[List[str], List[Tuple[str, str, str]]],
+        heads: int = 1,
+        negative_slope: float = 0.2,
+        **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.k_lin = keras.layers.Dense(out_channels, use_bias=True)
+        self.proj = {
+            node_type: keras.layers.Dense(out_channels, use_bias=True)
+            for node_type, ch in self.in_channels.items()
+        }
+
+    def build(self, input_shape=None):
+        H, D = self.heads, self.out_channels // self.heads
+
+        self.q = self.add_weight(
+            shape=(1, self.out_channels),
+            initializer="glorot_uniform",
+            trainable=True,
+            name="q",
+        )
+
+        self.lin_src = {}
+        self.lin_dst = {}
+        for edge_type in self.metadata[1]:
+            key = "__".join(edge_type)
+            self.lin_src[key] = self.add_weight(
+                shape=(1, H, D),
+                initializer="glorot_uniform",
+                trainable=True,
+                name=f"lin_src_{key}",
+            )
+            self.lin_dst[key] = self.add_weight(
+                shape=(1, H, D),
+                initializer="glorot_uniform",
+                trainable=True,
+                name=f"lin_dst_{key}",
+            )
+
+        for nt, layer in self.proj.items():
+            if not layer.built:
+                layer.build((None, self.in_channels[nt]))
+        if not self.k_lin.built:
+            self.k_lin.build((None, self.out_channels))
+
+        super().build(input_shape)
+
+    def call(
+        self,
+        x_dict: Dict[str, any],
+        edge_index_dict: Dict[Tuple[str, str, str], any],
+        return_semantic_attention_weights: bool = False,
+    ):
+        if not self.built:
+            self.build()
+
+        H, D = self.heads, self.out_channels // self.heads
+        x_node_dict = {}
+        out_dict = {nt: [] for nt in self.metadata[0]}
+
+        for node_type, x in x_dict.items():
+            proj_x = self.proj[node_type](x)
+            x_node_dict[node_type] = ops.reshape(proj_x, (-1, H, D))
+
+        for edge_type, edge_index in edge_index_dict.items():
+            src_type, _, dst_type = edge_type
+            key = "__".join(edge_type)
+            lin_src = self.lin_src[key]
+            lin_dst = self.lin_dst[key]
+            x_src = x_node_dict[src_type]
+            x_dst = x_node_dict[dst_type]
+
+            alpha_src = ops.sum(x_src * lin_src, axis=-1)  # (N_src, H)
+            alpha_dst = ops.sum(x_dst * lin_dst, axis=-1)  # (N_dst, H)
+
+            out = self.propagate(
+                edge_index,
+                x=(x_src, x_dst),
+                alpha=(alpha_src, alpha_dst),
+            )
+            out = ops.relu(out)
+            out_dict[dst_type].append(out)
+
+        result = {}
+        semantic_attn_dict = {}
+        for node_type, outs in out_dict.items():
+            out, attn = semantic_group(outs, self.q, self.k_lin)
+            result[node_type] = out
+            semantic_attn_dict[node_type] = attn
+
+        if return_semantic_attention_weights:
+            return result, semantic_attn_dict
+        return result
+
+    def message(self, x_j, alpha_i, alpha_j, index=None):
+        alpha = alpha_j + alpha_i
+        alpha = ops.leaky_relu(alpha, negative_slope=self.negative_slope)
+        alpha = softmax(alpha, index=index)
+        out = x_j * ops.expand_dims(alpha, axis=-1)
+        return ops.reshape(out, (-1, self.out_channels))
+

k3_node/layers/conv/heat_conv.py

@@ -0,0 +1,131 @@
+import keras
+from keras import ops
+from k3_node.layers.conv.message_passing import MessagePassing
+from k3_node.layers.conv.utils import softmax
+from k3_node.layers.dense.linear import HeteroLinear
+
+
+class HEATConv(MessagePassing):
+    r"""The heterogeneous edge-enhanced graph attentional operator from the
+    `"Heterogeneous Edge-Enhanced Graph Attention Network For Multi-Agent
+    Trajectory Prediction" <https://arxiv.org/abs/2106.07161>`_ paper.
+
+    Args:
+        in_channels (int): Size of each input sample.
+        out_channels (int): Size of each output sample.
+        num_node_types (int): The number of node types.
+        num_edge_types (int): The number of edge types.
+        edge_type_emb_dim (int): The embedding size of edge types.
+        edge_dim (int): Edge feature dimensionality.
+        edge_attr_emb_dim (int): The embedding size of edge features.
+        heads (int, optional): Number of multi-head-attentions. (default: :obj:`1`)
+        concat (bool, optional): Whether to concatenate multi-head attention. (default: :obj:`True`)
+        negative_slope (float, optional): LeakyReLU angle. (default: :obj:`0.2`)
+        root_weight (bool, optional): Whether to add root node features. (default: :obj:`True`)
+        bias (bool, optional): Whether to learn an additive bias. (default: :obj:`True`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import HEATConv
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges
+
+        node_type = np.random.randint(0, 2, size=(10,))
+        edge_type = np.random.randint(0, 3, size=(30,))
+        edge_attr = np.random.rand(30, 5).astype("float32")
+        layer = HEATConv(in_channels=8, out_channels=16, num_node_types=2, num_edge_types=3,
+                         edge_type_emb_dim=4, edge_dim=5, edge_attr_emb_dim=6, heads=2)
+        out = layer(x, edge_index, node_type, edge_type, edge_attr)
+        print(tuple(out.shape))  # (10, 32)
+        ```
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        num_node_types: int,
+        num_edge_types: int,
+        edge_type_emb_dim: int,
+        edge_dim: int,
+        edge_attr_emb_dim: int,
+        heads: int = 1,
+        concat: bool = True,
+        negative_slope: float = 0.2,
+        root_weight: bool = True,
+        bias: bool = True,
+        **kwargs,
+    ):
+        kwargs.setdefault("aggr", "add")
+        super().__init__(node_dim=0, **kwargs)
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.num_node_types = num_node_types
+        self.num_edge_types = num_edge_types
+        self.edge_type_emb_dim = edge_type_emb_dim
+        self.edge_dim = edge_dim
+        self.edge_attr_emb_dim = edge_attr_emb_dim
+        self.heads = heads
+        self.concat = concat
+        self.negative_slope = negative_slope
+        self.root_weight = root_weight
+        self.use_bias = bias
+
+        self.hetero_lin = HeteroLinear(in_channels, out_channels, num_node_types, bias=bias)
+        self.edge_type_emb = keras.layers.Embedding(num_edge_types, edge_type_emb_dim)
+        self.edge_attr_emb = keras.layers.Dense(edge_attr_emb_dim, use_bias=False)
+        self.att = keras.layers.Dense(heads, use_bias=False)
+        self.lin = keras.layers.Dense(out_channels, use_bias=bias)
+
+    def build(self, input_shape=None):
+        if not self.hetero_lin.built and self.in_channels > 0:
+            self.hetero_lin.build((None, self.in_channels))
+        if not self.edge_type_emb.built:
+            self.edge_type_emb.build((None,))
+        if not self.edge_attr_emb.built and self.edge_dim > 0:
+            self.edge_attr_emb.build((None, self.edge_dim))
+        att_dim = 2 * self.out_channels + self.edge_type_emb_dim + self.edge_attr_emb_dim
+        if not self.att.built:
+            self.att.build((None, att_dim))
+        lin_dim = self.out_channels + self.edge_attr_emb_dim
+        if not self.lin.built:
+            self.lin.build((None, lin_dim))
+        super().build(input_shape)
+
+    def call(self, x, edge_index, node_type, edge_type, edge_attr=None):
+        if not self.built:
+            self.build()
+
+        x = self.hetero_lin(x, node_type)
+        edge_type_emb = ops.leaky_relu(
+            self.edge_type_emb(edge_type),
+            negative_slope=self.negative_slope,
+        )
+
+        out = self.propagate(edge_index, x=x, edge_type_emb=edge_type_emb, edge_attr=edge_attr)
+
+        if self.concat:
+            if self.root_weight:
+                out = out + ops.expand_dims(x, axis=1)
+            out = ops.reshape(out, (-1, self.heads * self.out_channels))
+        else:
+            out = ops.mean(out, axis=1)
+            if self.root_weight:
+                out = out + x
+
+        return out
+
+    def message(self, x_i, x_j, edge_type_emb, edge_attr, index=None):
+        edge_attr = ops.leaky_relu(self.edge_attr_emb(edge_attr), negative_slope=self.negative_slope)
+        alpha = ops.concatenate([x_i, x_j, edge_type_emb, edge_attr], axis=-1)
+        alpha = ops.leaky_relu(self.att(alpha), negative_slope=self.negative_slope)
+        alpha = softmax(alpha, index=index)
+
+        feat = self.lin(ops.concatenate([x_j, edge_attr], axis=-1))
+        # feat: (E, out_channels), alpha: (E, heads)
+        out = ops.expand_dims(feat, axis=1) * ops.expand_dims(alpha, axis=-1)
+        return out
+

k3_node/layers/conv/hetero_conv.py

@@ -0,0 +1,128 @@
+from typing import Dict, List, Optional, Tuple
+import keras
+from keras import ops
+
+
+def group(xs: List[any], aggr: Optional[str]) -> Optional[any]:
+    if len(xs) == 0:
+        return None
+    elif aggr is None:
+        return ops.stack(xs, axis=1)
+    elif len(xs) == 1:
+        return xs[0]
+    elif aggr == "cat":
+        return ops.concatenate(xs, axis=-1)
+    elif aggr == "sum":
+        return ops.sum(ops.stack(xs, axis=0), axis=0)
+    elif aggr == "mean":
+        return ops.mean(ops.stack(xs, axis=0), axis=0)
+    elif aggr == "min":
+        return ops.min(ops.stack(xs, axis=0), axis=0)
+    elif aggr == "max":
+        return ops.max(ops.stack(xs, axis=0), axis=0)
+    else:
+        raise ValueError(f"Unsupported aggregation: {aggr}")
+
+
+class HeteroConv(keras.layers.Layer):
+    r"""A generic wrapper for computing graph convolution on heterogeneous graphs.
+
+    Args:
+        convs (Dict[Tuple[str, str, str], keras.layers.Layer]): A dictionary holding a
+            bipartite GNN layer for each individual edge type.
+        aggr (str, optional): The aggregation scheme to use for grouping node
+            embeddings generated by different relations (:obj:`"sum"`, :obj:`"mean"`,
+            :obj:`"min"`, :obj:`"max"`, :obj:`"cat"`, :obj:`None`). (default: :obj:`"sum"`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import HeteroConv, SAGEConv
+
+        x_dict = {
+            "author": np.random.rand(3, 8).astype("float32"),
+            "paper": np.random.rand(4, 8).astype("float32"),
+        }
+        edge_index_dict = {
+            ("author", "writes", "paper"): np.array([[0, 1, 2], [1, 2, 3]]),
+            ("paper", "written_by", "author"): np.array([[1, 2, 3], [0, 1, 2]]),
+        }
+
+        # One convolution per edge type; results for the same node type are summed
+        layer = HeteroConv({
+            ("author", "writes", "paper"): SAGEConv((8, 8), 16),
+            ("paper", "written_by", "author"): SAGEConv((8, 8), 16),
+        }, aggr="sum")
+        out_dict = layer(x_dict, edge_index_dict)
+        print(tuple(out_dict["author"].shape), tuple(out_dict["paper"].shape))  # (3, 16) (4, 16)
+        ```
+    """
+
+    def __init__(
+        self,
+        convs: Dict[Tuple[str, str, str], keras.layers.Layer],
+        aggr: Optional[str] = "sum",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.convs = convs
+        self.aggr = aggr
+
+    def build(self, input_shape=None):
+        self.built = True
+
+    def call(self, *args_dict, **kwargs_dict) -> Dict[str, any]:
+        out_dict: Dict[str, List[any]] = {}
+
+        for edge_type, conv in self.convs.items():
+            src, rel, dst = edge_type
+            has_edge_level_arg = False
+
+            args = []
+            for value_dict in args_dict:
+                if edge_type in value_dict:
+                    has_edge_level_arg = True
+                    args.append(value_dict[edge_type])
+                elif src == dst and src in value_dict:
+                    args.append(value_dict[src])
+                elif src in value_dict or dst in value_dict:
+                    args.append((
+                        value_dict.get(src, None),
+                        value_dict.get(dst, None),
+                    ))
+
+            kwargs = {}
+            for arg, value_dict in kwargs_dict.items():
+                if not arg.endswith("_dict"):
+                    raise ValueError(
+                        f"Keyword arguments in '{self.__class__.__name__}' "
+                        f"need to end with '_dict' (got '{arg}')"
+                    )
+                clean_arg = arg[:-5]
+                if edge_type in value_dict:
+                    has_edge_level_arg = True
+                    kwargs[clean_arg] = value_dict[edge_type]
+                elif src == dst and src in value_dict:
+                    kwargs[clean_arg] = value_dict[src]
+                elif src in value_dict or dst in value_dict:
+                    kwargs[clean_arg] = (
+                        value_dict.get(src, None),
+                        value_dict.get(dst, None),
+                    )
+
+            if not has_edge_level_arg:
+                continue
+
+            out = conv(*args, **kwargs)
+
+            if dst not in out_dict:
+                out_dict[dst] = [out]
+            else:
+                out_dict[dst].append(out)
+
+        result: Dict[str, any] = {}
+        for key, value in out_dict.items():
+            result[key] = group(value, self.aggr)
+
+        return result
+

k3_node/layers/conv/hgt_conv.py

@@ -0,0 +1,218 @@
+import math
+from typing import Dict, List, Tuple, Union
+from keras import ops
+from k3_node.layers.conv.message_passing import MessagePassing
+from k3_node.layers.conv.utils import softmax
+from k3_node.layers.dense.linear import HeteroDictLinear, HeteroLinear
+from k3_node.ops.creation import repeat
+
+
+class HGTConv(MessagePassing):
+    r"""The Heterogeneous Graph Transformer (HGT) operator from the
+    `"Heterogeneous Graph Transformer" <https://arxiv.org/abs/2003.01332>`_ paper.
+
+    Args:
+        in_channels (int or Dict[str, int]): Size of each input sample of every node type.
+        out_channels (int): Size of each output sample.
+        metadata (Tuple[List[str], List[Tuple[str, str, str]]]): Node types and edge types.
+        heads (int, optional): Number of multi-head-attentions. (default: :obj:`1`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import HGTConv
+
+        x_dict = {
+            "author": np.random.rand(3, 8).astype("float32"),
+            "paper": np.random.rand(4, 8).astype("float32"),
+        }
+        edge_index_dict = {
+            ("author", "writes", "paper"): np.array([[0, 1, 2], [1, 2, 3]]),
+            ("paper", "written_by", "author"): np.array([[1, 2, 3], [0, 1, 2]]),
+        }
+        metadata = (list(x_dict), list(edge_index_dict))  # (node types, edge types)
+
+        layer = HGTConv(in_channels={"author": 8, "paper": 8}, out_channels=16, metadata=metadata, heads=2)
+        out_dict = layer(x_dict, edge_index_dict)
+        print(tuple(out_dict["author"].shape), tuple(out_dict["paper"].shape))  # (3, 16) (4, 16)
+        ```
+    """
+
+    def __init__(
+        self,
+        in_channels: Union[int, Dict[str, int]],
+        out_channels: int,
+        metadata: Tuple[List[str], List[Tuple[str, str, str]]],
+        heads: int = 1,
+        **kwargs,
+    ):
+        super().__init__(aggr="add", node_dim=0, **kwargs)
+
+        if out_channels % heads != 0:
+            raise ValueError(f"'out_channels' ({out_channels}) must be divisible by 'heads' ({heads})")
+
+        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.node_types = metadata[0]
+        self.edge_types = metadata[1]
+        self.edge_types_map = {edge_type: i for i, edge_type in enumerate(metadata[1])}
+        self.dst_node_types = {key[-1] for key in self.edge_types}
+
+        self.kqv_lin = HeteroDictLinear(self.in_channels, self.out_channels * 3)
+        self.out_lin = HeteroDictLinear(self.out_channels, self.out_channels, types=self.node_types)
+
+        dim = out_channels // heads
+        num_types = heads * len(self.edge_types)
+
+        self.k_rel = HeteroLinear(dim, dim, num_types, bias=False, is_sorted=True)
+        self.v_rel = HeteroLinear(dim, dim, num_types, bias=False, is_sorted=True)
+
+    def build(self, input_shape=None):
+        self.skip = {}
+        for nt in self.node_types:
+            self.skip[nt] = self.add_weight(
+                shape=(1,),
+                initializer="ones",
+                trainable=True,
+                name=f"skip_{nt}",
+            )
+        self.p_rel = {}
+        for et in self.edge_types:
+            key = "__".join(et)
+            self.p_rel[key] = self.add_weight(
+                shape=(1, self.heads),
+                initializer="ones",
+                trainable=True,
+                name=f"p_rel_{key}",
+            )
+        super().build(input_shape)
+
+    def _cat(self, x_dict: Dict[str, any]) -> Tuple[any, Dict[str, int]]:
+        cumsum = 0
+        outs = []
+        offset = {}
+        for key, x in x_dict.items():
+            outs.append(x)
+            offset[key] = cumsum
+            cumsum += ops.shape(x)[0]
+        return ops.concatenate(outs, axis=0), offset
+
+    def _construct_src_node_feat(
+        self,
+        k_dict: Dict[str, any],
+        v_dict: Dict[str, any],
+        edge_index_dict: Dict[Tuple[str, str, str], any],
+    ):
+        cumsum = 0
+        num_edge_types = len(self.edge_types)
+        H, D = self.heads, self.out_channels // self.heads
+
+        ks = []
+        vs = []
+        type_list = []
+        offset = {}
+
+        for edge_type in edge_index_dict.keys():
+            src = edge_type[0]
+            N = ops.shape(k_dict[src])[0]
+            offset[edge_type] = cumsum
+            cumsum += N
+
+            edge_type_offset = self.edge_types_map[edge_type]
+            type_vec = (
+                repeat(ops.reshape(ops.arange(H, dtype="int32"), (-1, 1)), N, axis=1)
+                * num_edge_types
+                + edge_type_offset
+            )
+
+            type_list.append(type_vec)
+            ks.append(k_dict[src])
+            vs.append(v_dict[src])
+
+        ks_cat = ops.reshape(ops.transpose(ops.concatenate(ks, axis=0), (1, 0, 2)), (-1, D))
+        vs_cat = ops.reshape(ops.transpose(ops.concatenate(vs, axis=0), (1, 0, 2)), (-1, D))
+        type_vec_cat = ops.reshape(ops.concatenate(type_list, axis=1), (-1,))
+
+        k = self.k_rel(ks_cat, type_vec_cat)
+        k = ops.transpose(ops.reshape(k, (H, -1, D)), (1, 0, 2))
+
+        v = self.v_rel(vs_cat, type_vec_cat)
+        v = ops.transpose(ops.reshape(v, (H, -1, D)), (1, 0, 2))
+
+        return k, v, offset
+
+    def call(
+        self,
+        x_dict: Dict[str, any],
+        edge_index_dict: Dict[Tuple[str, str, str], any],
+    ) -> Dict[str, any]:
+        if not self.built:
+            self.build()
+
+        F = self.out_channels
+        H = self.heads
+        D = F // H
+
+        k_dict, q_dict, v_dict, out_dict = {}, {}, {}, {}
+
+        kqv_dict = self.kqv_lin(x_dict)
+        for key, val in kqv_dict.items():
+            k, q, v = ops.split(val, 3, axis=1)
+            k_dict[key] = ops.reshape(k, (-1, H, D))
+            q_dict[key] = ops.reshape(q, (-1, H, D))
+            v_dict[key] = ops.reshape(v, (-1, H, D))
+
+        q, dst_offset = self._cat(q_dict)
+        k, v, src_offset = self._construct_src_node_feat(k_dict, v_dict, edge_index_dict)
+
+        # Build concatenated bipartite edge index and edge attributes
+        edge_indices = []
+        edge_attrs = []
+        for edge_type, s_offset in src_offset.items():
+            e_idx = edge_index_dict[edge_type]
+            d_offset = dst_offset[edge_type[-1]]
+
+            row = e_idx[0] + s_offset
+            col = e_idx[1] + d_offset
+            edge_indices.append(ops.stack([row, col], axis=0))
+
+            p_val = self.p_rel["__".join(edge_type)]
+            num_e = ops.shape(e_idx)[1]
+            edge_attrs.append(repeat(p_val, num_e, axis=0))
+
+        edge_index = ops.concatenate(edge_indices, axis=1)
+        edge_attr = ops.concatenate(edge_attrs, axis=0)
+
+        out = self.propagate(edge_index, k=k, q=q, v=v, edge_attr=edge_attr)
+
+        for node_type, start_offset in dst_offset.items():
+            count = ops.shape(q_dict[node_type])[0]
+            if node_type in self.dst_node_types:
+                out_dict[node_type] = out[start_offset : start_offset + count]
+
+        transformed = {}
+        for k_nt, v_nt in out_dict.items():
+            transformed[k_nt] = ops.gelu(v_nt) if v_nt is not None else v_nt
+        a_dict = self.out_lin(transformed)
+
+        result = {}
+        for node_type, o in out_dict.items():
+            res = a_dict[node_type]
+            if ops.shape(res)[-1] == ops.shape(x_dict[node_type])[-1]:
+                alpha = ops.sigmoid(self.skip[node_type])
+                res = alpha * res + (1.0 - alpha) * x_dict[node_type]
+            result[node_type] = res
+
+        return result
+
+    def message(self, k_j, q_i, v_j, edge_attr, index=None):
+        alpha = ops.sum(q_i * k_j, axis=-1) * edge_attr
+        alpha = alpha / math.sqrt(ops.shape(q_i)[-1])
+        alpha = softmax(alpha, index=index)
+        out = v_j * ops.expand_dims(alpha, axis=-1)
+        return ops.reshape(out, (-1, self.out_channels))
+