def METHOD_NAME(
    self,
    x_dict: Dict[NodeType, Tensor],
    edge_index_dict: Dict[EdgeType, Adj],
    return_semantic_attention_weights: bool = False,
) -> Union[Dict[NodeType, OptTensor], Tuple[Dict[NodeType, OptTensor],
                                            Dict[NodeType, OptTensor]]]:
    r"""Runs the forward pass of the module.
    Args:
        x_dict (Dict[str, torch.Tensor]): A dictionary holding node feature
            information for each individual node type.
        edge_index_dict (Dict[Tuple[str, str, str], torch.Tensor]): A
            dictionary holding graph connectivity information for each
            individual edge type, either as a :class:`torch.Tensor` of
            shape :obj:`[2, num_edges]` or a
            :class:`torch_sparse.SparseTensor`.
        return_semantic_attention_weights (bool, optional): If set to
            :obj:`True`, will additionally return the semantic-level
            attention weights for each destination node type.
            (default: :obj:`False`)
    """
    H, D = self.heads, self.out_channels // self.heads
    x_node_dict, out_dict = {}, {}
    # Iterate over node types:
    for node_type, x in x_dict.items():
        x_node_dict[node_type] = self.proj[node_type](x).view(-1, H, D)
        out_dict[node_type] = []
    # Iterate over edge types:
    for edge_type, edge_index in edge_index_dict.items():
        src_type, _, dst_type = edge_type
        edge_type = '__'.join(edge_type)
        lin_src = self.lin_src[edge_type]
        lin_dst = self.lin_dst[edge_type]
        x_src = x_node_dict[src_type]
        x_dst = x_node_dict[dst_type]
        alpha_src = (x_src * lin_src).sum(dim=-1)
        alpha_dst = (x_dst * lin_dst).sum(dim=-1)
        # propagate_type: (x_dst: PairTensor, alpha: PairTensor)
        out = self.propagate(edge_index, x=(x_src, x_dst),
                             alpha=(alpha_src, alpha_dst), size=None)
        out = F.relu(out)
        out_dict[dst_type].append(out)
    # iterate over node types:
    semantic_attn_dict = {}
    for node_type, outs in out_dict.items():
        out, attn = group(outs, self.q, self.k_lin)
        out_dict[node_type] = out
        semantic_attn_dict[node_type] = attn
    if return_semantic_attention_weights:
        return out_dict, semantic_attn_dict
    return out_dict