Models¶

class CommonNeighbors(nn.Module):
    def __init__(
        self,
        aggregation: Literal["mean", "min", "sum"],
        scorer: NeighborScorer | None = None,
    ) -> None:
        super().__init__()
        self.scorer = scorer if scorer is not None else CommonNeighborsScorer(aggregation)

    def forward(
        self,
        hyperedge_index: Tensor,
        node_to_neighbors: dict[int, Neighborhood] | None = None,
    ) -> Tensor:
        """
        Compute CN scores for all hyperedges in the batch.

        Args:
            hyperedge_index: Tensor containing the hyperedge indices.
            node_to_neighbors: Optional mapping from nodes to their neighborhoods.

        Returns:
            A 1-D tensor of shape (num_hyperedges,) with CN scores.
        """
        scores = self.scorer.score_batch(hyperedge_index, node_to_neighbors)
        torch.log1p(scores, out=scores)
        return scores

class GCN(nn.Module):
    """
    A reusable multi-layer GCN stack built from ``torch_geometric.nn.GCNConv``.

    Args:
        in_channels: Dimension of the input node embeddings to the GCN layers.
        out_channels: Dimension of the output node embeddings from the GCN layers.
        hidden_channels: Dimension of the hidden node embeddings in the GCN layers.
            Defaults to ``in_channels``.
        num_layers: Number of GCN layers. Must be at least 1. Defaults to ``2``.
        drop_rate: Dropout rate applied after each GCN layer except the last one.
        bias: Whether to include a bias term in the GCN layers.
        activation_fn: Activation function to use after each hidden layer. Defaults to ``nn.ReLU``.
        activation_fn_kwargs: Keyword arguments for the activation function. Defaults to empty dict.
        improved: Whether to use the improved version of ``GCNConv``.
        add_self_loops: Whether to add self-loops to the input graph.
        normalize: Whether to symmetrically normalize the adjacency matrix in ``GCNConv``.
        cached: Whether to cache the normalized adjacency matrix in ``GCNConv``.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        hidden_channels: int | None = None,
        num_layers: int = 2,
        drop_rate: float = 0.0,
        bias: bool = True,
        activation_fn: ActivationFn | None = None,
        activation_fn_kwargs: dict | None = None,
        improved: bool = False,
        add_self_loops: bool = True,
        normalize: bool = True,
        cached: bool = False,
    ):
        super().__init__()
        activation_fn = activation_fn if activation_fn is not None else nn.ReLU
        activation_fn_kwargs = activation_fn_kwargs if activation_fn_kwargs is not None else {}

        self.dropout = nn.Dropout(drop_rate)
        self.activation = activation_fn(**activation_fn_kwargs)
        self.layers = self.__build_layers(
            in_channels=in_channels,
            out_channels=out_channels,
            hidden_channels=hidden_channels,
            num_layers=num_layers,
            bias=bias,
            improved=improved,
            add_self_loops=add_self_loops,
            normalize=normalize,
            cached=cached,
        )

    def forward(self, x: Tensor, edge_index: Tensor) -> Tensor:
        num_layers = len(self.layers)
        for idx, layer in enumerate(self.layers):
            x = layer(x, edge_index)

            is_not_last_layer = not is_layer(idx, num_layers - 1)
            if is_not_last_layer:
                x = self.activation(x)
                x = self.dropout(x)

        return x

    def __build_layers(
        self,
        in_channels: int,
        out_channels: int,
        hidden_channels: int | None,
        num_layers: int,
        bias: bool,
        improved: bool,
        add_self_loops: bool,
        normalize: bool,
        cached: bool,
    ) -> nn.ModuleList:
        if num_layers < 1:
            raise ValueError(f"Expected num_layers >= 1 for GCN, got {num_layers}.")

        hidden_channels = hidden_channels if hidden_channels is not None else 0
        if num_layers > 1 and hidden_channels <= 0:
            raise ValueError(
                f"Expected positive hidden_channels for GCN with multiple layers, got {hidden_channels}."
            )

        common_kwargs: dict[str, bool] = {
            "bias": bias,
            "improved": improved,
            "add_self_loops": add_self_loops,
            "normalize": normalize,
            "cached": cached,
        }

        if num_layers == 1:
            return nn.ModuleList([GCNConv(in_channels, out_channels, **common_kwargs)])

        layers = [GCNConv(in_channels, hidden_channels, **common_kwargs)]
        layers.extend(
            GCNConv(hidden_channels, hidden_channels, **common_kwargs)
            for _ in range(num_layers - 2)
        )
        layers.append(GCNConv(hidden_channels, out_channels, **common_kwargs))

        return nn.ModuleList(layers)

class GCNConfig(TypedDict):
    """
    Configuration for the GCN model.

    Args:
        in_channels: Dimension of the input node embeddings to the GCN layers.
        out_channels: Dimension of the output node embeddings from the GCN layers.
        hidden_channels: Dimension of the hidden node embeddings in the GCN layers.
        num_layers: Number of GCN layers. Must be at least 1. Defaults to ``2``.
        drop_rate: Dropout rate applied after each GCN layer (except the last one). Defaults to ``0.0`` (no dropout).
        activation_fn: Activation function to use after each hidden layer. Defaults to ``nn.ReLU``.
        activation_fn_kwargs: Keyword arguments for the activation function. Defaults to empty dict.
        bias: Whether to include a bias term in the GCN layers. Defaults to ``True``.
        improved: Whether to use the improved version of GCNConv. Defaults to ``False``.
        add_self_loops: Whether to add self-loops to the input graph. Defaults to ``True``.
        normalize: Whether to symmetrically normalize the adjacency matrix in GCNConv. Defaults to ``True``.
        cached: Whether to cache the normalized adjacency matrix in GCNConv.
            Only applicable if the graph structure does not change between epochs. Defaults to ``False``.
    """

    in_channels: int
    out_channels: int
    hidden_channels: NotRequired[int]
    num_layers: NotRequired[int]
    drop_rate: NotRequired[float]
    bias: NotRequired[bool]
    activation_fn: NotRequired[ActivationFn]
    activation_fn_kwargs: NotRequired[dict]
    improved: NotRequired[bool]
    add_self_loops: NotRequired[bool]
    normalize: NotRequired[bool]
    cached: NotRequired[bool]

class HGNN(nn.Module):
    """
    HGNN performs spectral convolution directly on the hypergraph structure using the
    node-hyperedge incidence matrix, without any reduction to a pairwise graph.
    Unlike HyperGCN (which approximates each hyperedge by selecting representative pairwise
    edges via random projection), HGNN preserves all higher-order relationships by passing
    messages through the full incidence structure: nodes -> hyperedges -> nodes.
    - Proposed in `Hypergraph Neural Networks <https://arxiv.org/pdf/1809.09401>`_ paper (AAAI 2019).
    - Reference implementation: `source <https://deephypergraph.readthedocs.io/en/latest/_modules/dhg/models/hypergraphs/hgnn.html#HGNN>`_.

    Args:
        in_channels: The number of input channels.
        hidden_channels: The number of hidden channels.
        num_classes: The number of output channels.
        bias: If set to ``False``, the layer will not learn the bias parameter. Defaults to ``True``.
        use_batch_normalization: If set to ``True``, layers will use batch normalization. Defaults to ``False``.
        drop_rate: Dropout ratio. Defaults to ``0.5``.
    """

    def __init__(
        self,
        in_channels: int,
        hidden_channels: int,
        num_classes: int,
        bias: bool = True,
        use_batch_normalization: bool = False,
        drop_rate: float = 0.5,
    ):
        super().__init__()

        self.layers = nn.ModuleList(
            [
                HGNNConv(
                    in_channels=in_channels,
                    out_channels=hidden_channels,
                    bias=bias,
                    use_batch_normalization=use_batch_normalization,
                    drop_rate=drop_rate,
                ),
                HGNNConv(
                    in_channels=hidden_channels,
                    out_channels=num_classes,
                    bias=bias,
                    use_batch_normalization=use_batch_normalization,
                    is_last=True,
                ),
            ]
        )

    def forward(self, x: Tensor, hyperedge_index: Tensor) -> Tensor:
        """
        Apply two stacked ``HGNNConv`` layers to produce node embeddings.

        The first layer applies ReLU + dropout and maps ``in_channels -> hidden_channels``.
        The second layer is the output layer (no activation/dropout) and maps
        ``hidden_channels -> num_classes``.

        Args:
            x: Input node feature matrix. Size ``(num_nodes, in_channels)``.
            hyperedge_index: Hyperedge incidence in COO format. Size ``(2, num_incidences)``,
                where row 0 contains node IDs and row 1 contains hyperedge IDs.

        Returns:
            The output node feature matrix. Size ``(num_nodes, num_classes)``.
        """
        for layer in self.layers:
            x = layer(x, hyperedge_index)
        return x

class HNHN(nn.Module):
    """
    HNHN performs incidence-based hypergraph convolution with explicit hyperedge
    embeddings between the node -> hyperedge -> node propagation steps.
    - Proposed in `HNHN: Hypergraph Networks with Hyperedge Neurons <https://arxiv.org/abs/2006.12278>`_ paper.
    - Reference implementation: `source <https://deephypergraph.readthedocs.io/en/latest/_modules/dhg/models/hypergraphs/hnhn.html#HNHN>`_.

    Args:
        in_channels: The number of input channels.
        hidden_channels: The number of hidden channels.
        num_classes: The number of output channels.
        bias: If set to ``False``, the layer will not learn the bias parameter. Defaults to ``True``.
        use_batch_normalization: If set to ``True``, layers will use batch normalization. Defaults to ``False``.
        drop_rate: Dropout ratio. Defaults to ``0.5``.
    """

    def __init__(
        self,
        in_channels: int,
        hidden_channels: int,
        num_classes: int,
        bias: bool = True,
        use_batch_normalization: bool = False,
        drop_rate: float = 0.5,
    ):
        super().__init__()

        self.layers = nn.ModuleList(
            [
                HNHNConv(
                    in_channels=in_channels,
                    out_channels=hidden_channels,
                    bias=bias,
                    use_batch_normalization=use_batch_normalization,
                    drop_rate=drop_rate,
                ),
                HNHNConv(
                    in_channels=hidden_channels,
                    out_channels=num_classes,
                    bias=bias,
                    use_batch_normalization=use_batch_normalization,
                    is_last=True,
                ),
            ]
        )

    def forward(self, x: Tensor, hyperedge_index: Tensor) -> Tensor:
        """
        Apply two stacked ``HNHNConv`` layers to produce node embeddings.

        Args:
            x: Input node feature matrix of size ``(num_nodes, in_channels)``.
            hyperedge_index: Hyperedge incidence in COO format of size ``(2, num_incidences)``.

        Returns:
            The output node feature matrix of size ``(num_nodes, num_classes)``.
        """
        for layer in self.layers:
            x = layer(x, hyperedge_index)
        return x

class HGNNP(nn.Module):
    """
    HGNN+ performs hypergraph convolution with two-stage mean aggregation using the
    incidence structure directly: nodes -> hyperedges -> nodes.
    - Proposed in `HGNN+: General Hypergraph Neural Networks <https://ieeexplore.ieee.org/document/9795251>`_ paper (IEEE T-PAMI 2022).
    - Reference implementation: `source <https://deephypergraph.readthedocs.io/en/latest/_modules/dhg/models/hypergraphs/hgnnp.html#HGNNP>`_.

    Args:
        in_channels: The number of input channels.
        hidden_channels: The number of hidden channels.
        num_classes: The number of output channels.
        bias: If set to ``False``, the layer will not learn the bias parameter. Defaults to ``True``.
        use_batch_normalization: If set to ``True``, layers will use batch normalization. Defaults to ``False``.
        drop_rate: Dropout ratio. Defaults to ``0.5``.
    """

    def __init__(
        self,
        in_channels: int,
        hidden_channels: int,
        num_classes: int,
        bias: bool = True,
        use_batch_normalization: bool = False,
        drop_rate: float = 0.5,
    ):
        super().__init__()

        self.layers = nn.ModuleList(
            [
                HGNNPConv(
                    in_channels=in_channels,
                    out_channels=hidden_channels,
                    bias=bias,
                    use_batch_normalization=use_batch_normalization,
                    drop_rate=drop_rate,
                ),
                HGNNPConv(
                    in_channels=hidden_channels,
                    out_channels=num_classes,
                    bias=bias,
                    use_batch_normalization=use_batch_normalization,
                    is_last=True,
                ),
            ]
        )

    def forward(self, x: Tensor, hyperedge_index: Tensor) -> Tensor:
        """
        Apply two stacked ``HGNNPConv`` layers to produce node embeddings.

        Args:
            x: Input node feature matrix. Size ``(num_nodes, in_channels)``.
            hyperedge_index: Hyperedge incidence in COO format. Size ``(2, num_incidences)``,
                where row 0 contains node IDs and row 1 contains hyperedge IDs.

        Returns:
            The output node feature matrix. Size ``(num_nodes, num_classes)``.
        """
        for layer in self.layers:
            x = layer(x, hyperedge_index)
        return x

class HyperGCN(nn.Module):
    """
    HyperGCN approximates each hyperedge of the hypergraph by a set of pairwise edges connecting the vertices of the hyperedge
    and treats the learning problem as a graph learning problem on the approximation.
    - Proposed in `HyperGCN: A New Method of Training Graph Convolutional Networks on Hypergraphs <https://dl.acm.org/doi/10.5555/3454287.3454422>`_ paper (NeurIPS 2019).
    - Code of the paper: `source <https://github.com/malllabiisc/HyperGCN>`_.
    - Reference implementation: `source <https://deephypergraph.readthedocs.io/en/latest/_modules/dhg/models/hypergraphs/hypergcn.html#HyperGCN>`_.

    Args:
        in_channels: The number of input channels.
        hidden_channels: The number of hidden channels.
        num_classes: The number of classes of the classification task as HyperGCB is a node classification model.
        bias: If set to ``False``, the layer will not learn the bias parameter. Defaults to ``True``.
        use_batch_normalization: If set to ``True``, layers will use batch normalization. Defaults to ``False``.
        drop_rate: Dropout ratio. Defaults to ``0.5``.
        use_mediator: Whether to use mediator to transform the hyperedges to edges in the graph. Defaults to ``False``.
        fast: If set to ``True``, the transformed graph structure will be computed once from the input hypergraph and vertex features, and cached for future use. Defaults to ``True``.
    """

    def __init__(
        self,
        in_channels: int,
        hidden_channels: int,
        num_classes: int,
        bias: bool = True,
        use_batch_normalization: bool = False,
        drop_rate: float = 0.5,
        use_mediator: bool = False,
        fast: bool = True,
    ):
        super().__init__()
        self.fast = fast
        self.use_mediator = use_mediator
        self.cached_gcn_laplacian_matrix: Tensor | None = None

        self.layers = nn.ModuleList(
            [
                HyperGCNConv(
                    in_channels=in_channels,
                    out_channels=hidden_channels,
                    bias=bias,
                    use_batch_normalization=use_batch_normalization,
                    drop_rate=drop_rate,
                    use_mediator=use_mediator,
                ),
                HyperGCNConv(
                    in_channels=hidden_channels,
                    out_channels=num_classes,
                    bias=bias,
                    use_batch_normalization=use_batch_normalization,
                    use_mediator=use_mediator,
                    is_last=True,
                ),
            ]
        )

    def forward(self, x: Tensor, hyperedge_index: Tensor) -> Tensor:
        """
        The forward function.

        Args:
            x: Input node feature matrix. Size ``(num_nodes, in_channels)``.
            hyperedge_index: The hyperedge indices of the hypergraph. Size ``(2, num_hyperedges)``.

        Returns:
            The output node feature matrix. Size ``(num_nodes, num_classes)``.
        """
        if not self.fast:
            for layer in self.layers:
                x = layer(x, hyperedge_index)
            return x

        # If the GCN Laplacian is cached, we need to check if the node feature size has changed
        # with cached_gcn_laplacian_matrix.size(0) != x.size(0), this can happen, for example, due to:
        # adding new negative samples or having validation/test sets with different node features
        should_not_use_cached_gcn_laplacian_matrix = (
            self.cached_gcn_laplacian_matrix is None  # Not cached yet
            or self.cached_gcn_laplacian_matrix.size(0)
            != x.size(0)  # Node feature size has changed
        )

        if should_not_use_cached_gcn_laplacian_matrix:
            edge_index, edge_weights = HyperedgeIndex(
                hyperedge_index
            ).reduce_to_edge_index_on_random_direction(
                x=x,
                with_mediators=self.use_mediator,
                return_weights=True,
            )

            self.cached_gcn_laplacian_matrix = EdgeIndex(
                edge_index=edge_index,
                edge_weights=edge_weights,
            ).get_sparse_normalized_gcn_laplacian(num_nodes=x.size(0))

        for layer in self.layers:
            x = layer(x, hyperedge_index, gcn_laplacian_matrix=self.cached_gcn_laplacian_matrix)
        return x

class MLP(nn.Module):
    """
    A simple multi-layer perceptron (MLP) with configurable number of layers, hidden channels, activation functions, normalization, and dropout.

    Examples:
        >>> mlp = MLP(in_channels=16, out_channels=1, hidden_channels=32, num_layers=3)
        >>> x = torch.randn(10, 16)  # 10 samples, 16 features
        >>> output = mlp(x)
        >>> output.shape
        ... torch.Size([10, 1])

        With custom activation, normalization, and dropout:
        >>> mlp = MLP(
        ...     in_channels=16,
        ...     out_channels=1,
        ...     hidden_channels=32,
        ...     num_layers=3,
        ...     activation_fn=nn.Tanh,                   # nn.ReLU, nn.LeakyReLU, etc.
        ...     activation_fn_kwargs={"inplace": True},
        ...     normalization_fn=nn.BatchNorm1d,         # nn.LayerNorm, etc.
        ...     normalization_fn_kwargs={"eps": 1e-5},
        ...     drop_rate=0.5,
        ... )
        >>> x = torch.randn(10, 16)
        >>> output = mlp(x)
        >>> output.shape
        ... torch.Size([10, 1])

    Args:
        in_channels: Number of input features.
        out_channels: Number of output features.
        hidden_channels: Number of hidden units in each hidden layer. Required if num_layers > 1.
        num_layers: Total number of layers (including output layer). Must be at least 1. Defaults to 1.
        activation_fn: Activation function to use after each hidden layer. Defaults to ``nn.ReLU``.
        activation_fn_kwargs: Keyword arguments for the activation function. Defaults to empty dict.
        normalization_fn: Normalization function to use after each hidden layer (before activation). If ``None``, no normalization is applied. Defaults to ``None``.
        normalization_fn_kwargs: Keyword arguments for the normalization function. Defaults to empty dict.
        bias: Whether to include bias terms in the linear layers. Defaults to ``True``.
        drop_rate: Dropout rate to apply after each hidden layer (after activation). If 0.0, no dropout is applied. Defaults to 0.0.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        hidden_channels: int | None = None,
        num_layers: int = 1,
        activation_fn: ActivationFn | None = None,
        activation_fn_kwargs: dict | None = None,
        normalization_fn: NormalizationFn | None = None,
        normalization_fn_kwargs: dict | None = None,
        bias: bool = True,
        drop_rate: float = 0.0,
    ):
        super().__init__()
        self.__validate_num_layers(num_layers, hidden_channels)

        hidden_channels = hidden_channels if hidden_channels is not None else 0
        activation_fn = activation_fn if activation_fn is not None else nn.ReLU
        activation_fn_kwargs = activation_fn_kwargs if activation_fn_kwargs is not None else {}
        normalization_fn_kwargs = (
            normalization_fn_kwargs if normalization_fn_kwargs is not None else {}
        )

        layers = nn.ModuleList()
        for layer_idx in range(num_layers):
            is_output_layer = is_layer(layer_idx, num_layers - 1)

            linear_layer = nn.Linear(
                in_features=in_channels if is_input_layer(layer_idx) else hidden_channels,
                out_features=out_channels if is_output_layer else hidden_channels,
                bias=bias,
            )
            layers.append(linear_layer)

            if not is_output_layer:
                if normalization_fn is not None:
                    layers.append(normalization_fn(hidden_channels, **normalization_fn_kwargs))

                layers.append(activation_fn(**activation_fn_kwargs))

                if drop_rate > 0.0:
                    layers.append(nn.Dropout(drop_rate))

        self.layers = nn.Sequential(*layers)

    def forward(self, x) -> Tensor:
        return self.layers(x)

    def __validate_num_layers(self, num_layers: int, hidden_channels: int | None) -> None:
        if num_layers < 1:
            raise ValueError("At least one layer is required for MLP.")
        if num_layers > 1 and hidden_channels is None:
            raise ValueError("hidden_channels must be specified for MLP with more than 1 layer.")

class SLP(MLP):
    """
    A single-layer perceptron (SLP) which is a special case of MLP with exactly one layer and no hidden units.

    Examples:
        >>> slp = SLP(in_channels=16, out_channels=1)
        >>> x = torch.randn(10, 16)  # 10 samples, 16 features
        >>> output = slp(x)
        >>> output.shape
        ... torch.Size([10, 1])

    Args:
        in_channels: Number of input features.
        out_channels: Number of output features.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
    ):
        super().__init__(
            in_channels=in_channels,
            out_channels=out_channels,
            num_layers=1,
        )

class NHP(nn.Module):
    """
    Neural Hyperlink Predictor (NHP) for undirected hyperedge link prediction.
    - Proposed in `NHP: Neural Hypergraph Link Prediction <https://dl.acm.org/doi/10.1145/3340531.3411870>`_ paper (CIKM 2020).
    - Reference implementation: `source <https://github.com/cyixiao/NHP-reproduce/>`_.

    NHP scores each candidate hyperedge by building candidate-specific node embeddings.
    A node that appears in multiple candidate hyperedges can receive a different incidence embedding in each one,
    because its update depends on the other nodes in that candidate hyperedge.

    Examples:
        >>> x = [
        ...     [1., 0.],  # node 0
        ...     [0., 1.],  # node 1
        ...     [1., 1.],  # node 2
        ...     [1., 0.],  # node 3
        ... ]
        >>> hyperedge_index = [
        ...     [0, 1, 1, 2, 3],  # node IDs
        ...     [0, 0, 1, 1, 1],  # hyperedge IDs
        ... ]
        >>> # hyperedge 0 = {node 0, node 1}
        >>> # hyperedge 1 = {node 1, node 2, node 3}
        >>> model = NHP(in_channels=2, hidden_channels=8, aggregation="maxmin")
        >>> scores = model(x, hyperedge_index)
        >>> scores.shape
        ... torch.Size([2])

    Args:
        in_channels: Number of input features per node.
        hidden_channels: Number of hidden units in the node embeddings.
        activation_fn: Activation function to use after the linear transformations. Defaults to ``nn.ReLU``.
        activation_fn_kwargs: Keyword arguments for the activation function. Defaults to empty dict.
        aggregation: Method to aggregate the incidence embeddings into a hyperedge embedding. Must be either "maxmin" or "mean". Defaults to "maxmin".
        bias: Whether to include bias terms in the linear layers. Defaults to ``True``.
    """

    def __init__(
        self,
        in_channels: int,
        hidden_channels: int,
        activation_fn: ActivationFn | None = None,
        activation_fn_kwargs: dict | None = None,
        aggregation: Literal["mean", "maxmin"] = "maxmin",
        bias: bool = True,
    ):
        super().__init__()

        activation_fn = activation_fn if activation_fn is not None else nn.ReLU
        activation_fn_kwargs = activation_fn_kwargs if activation_fn_kwargs is not None else {}

        self.aggregation = aggregation

        self.self_loop = nn.Linear(in_channels, hidden_channels, bias=bias)
        # GCN message passing is implemented through neighbor sum computation,
        # so one projection is enough for the hyperedge-aware term
        self.hyperedge_aware = nn.Linear(in_channels, hidden_channels, bias=bias)
        self.activation_fn = activation_fn(**activation_fn_kwargs)

        self.hyperedge_score = nn.Linear(hidden_channels, 1, bias=bias)

    def forward(self, x: Tensor, hyperedge_index: Tensor) -> Tensor:
        """
        Score each candidate hyperedge.

        Args:
            x: Node feature matrix of shape ``(num_nodes, in_channels)``.
            hyperedge_index: Incidence tensor of shape ``(2, num_incidences)``.

        Returns:
            Scores of shape ``(num_hyperedges,)``.
        """
        if hyperedge_index.numel() == 0:
            return x.new_empty((0,))

        # Example: hyperedge_index = [[0, 1, 1, 2, 3],  == node_ids
        #                             [0, 0, 1, 1, 1]]  == hyperedge_ids
        node_ids = hyperedge_index[0]
        hyperedge_ids = hyperedge_index[1]

        # Gather the node features for each incidence
        # Example: x = [[1, 0],  # node 0
        #               [0, 1],  # node 1
        #               [1, 1],  # node 2
        #               [1, 0]]  # node 3
        #          node_ids = [0, 1, 1, 2, 3]
        #          -> incidence_node_features = [[1, 0],  # node 0 in hyperedge 0
        #                                        [0, 1],  # node 1 in hyperedge 0
        #                                        [0, 1],  # node 1 in hyperedge 1
        #                                        [1, 1],  # node 2 in hyperedge 1
        #                                        [1, 0]]  # node 3 in hyperedge 1
        #             shape: (num_incidences, in_channels)
        incidence_node_features = x[node_ids]

        # Do one local message-passing step to sum original node features per hyperedge to get hyperedge features.
        # that are aware of all nodes in the candidate hyperedge.
        # Example: hyperedge 0 contains nodes (0, 1)    -> [1, 0] + [0, 1] = [1, 1]
        #          hyperedge 1 contains nodes (1, 2, 3) -> [0, 1] + [1, 1] + [1, 0] = [2, 2]
        #          -> hyperedge_features = [[1, 1],  # sum for hyperedge 0
        #                                   [2, 2]]  # sum for hyperedge 1
        #             shape: (num_hyperedges, in_channels)
        hyperedge_features = HyperedgeAggregator(
            hyperedge_index=hyperedge_index,
            node_embeddings=x,
        ).pool("sum")

        # Broadcast hyperedge features back to each of their incidences,
        # and remove the current node feature to give to each incidence
        # the features of its neighboring nodes in the candidate hyperedge.
        # Example: hyperedge_features = [[1, 1],  # sum for hyperedge 0
        #                                [2, 2]]  # sum for hyperedge 1
        #                               shape (num_hyperedges, in_channels),
        #          hyperedge_ids = [0, 0, 1, 1, 1],
        #          incidence_node_features = [[1, 0],  # node 0 in hyperedge 0
        #                                     [0, 1],  # node 1 in hyperedge 0
        #                                     [0, 1],  # node 1 in hyperedge 1
        #                                     [1, 1],  # node 2 in hyperedge 1
        #                                     [1, 0]]  # node 3 in hyperedge 1
        #                                    shape: (num_incidences, in_channels)
        #          -> hyperedge_features[hyperedge_ids] = [[1, 1],  # hyperedge 0 for node 0
        #                                                  [1, 1],  # hyperedge 0 for node 1
        #                                                  [2, 2],  # hyperedge 1 for node 1
        #                                                  [2, 2],  # hyperedge 1 for node 2
        #                                                  [2, 2]]  # hyperedge 1 for node 3
        #                                                 shape: (num_incidences, in_channels)
        #          -> neighbor_features_per_incidence = [[0, 1],  # node 0 sees node 1
        #                                                [1, 0],  # node 1 sees node 0
        #                                                [2, 1],  # node 1 sees node 2 and node 3
        #                                                [1, 1],  # node 2 sees node 1 and node 3
        #                                                [1, 2]]  # node 3 sees node 1 and node 2
        #                                               shape: (num_incidences, in_channels)
        neighbor_features_per_incidence = (
            hyperedge_features[hyperedge_ids] - incidence_node_features
        )

        # shape (num_incidences, hidden_channels)
        neighbor_aware_hyperedge_embeddings = self.hyperedge_aware(neighbor_features_per_incidence)
        # shape (num_incidences, hidden_channels)
        selfloop_embeddings = self.self_loop(incidence_node_features)

        # incidence_embeddings[0] = activation_fn(selfloop_embeddings[0] + neighbor_aware_hyperedge_embeddings[0])
        # is the embedding of the first incidence (i.e., node 0 in hyperedge 0)
        # after one local message-passing step inside that candidate hyperedge.
        incidence_embeddings = self.activation_fn(
            selfloop_embeddings + neighbor_aware_hyperedge_embeddings
        )  # shape (num_incidences, hidden_channels)

        # Treat each incidence embedding as a separately aggregatable set of features.
        # This is required because incidence embeddings are not global node embeddings:
        # node 1 may appear twice with two different embeddings as it participates in two different candidate hyperedges.
        # Example: incidence_ids = [0, 1, 2, 3, 4],
        #          hyperedge_ids = [0, 0, 1, 1, 1]
        #          -> incidence_hyperedge_index = [[0, 1, 2, 3, 4],
        #                                          [0, 0, 1, 1, 1]]
        num_incidences = incidence_embeddings.size(0)
        incidence_ids = torch.arange(num_incidences, device=hyperedge_index.device)
        incidence_hyperedge_index = torch.stack([incidence_ids, hyperedge_ids], dim=0)

        # Example: incidence_embeddings = [[1, 2],  # features 0, node 0 in hyperedge 0
        #                                  [3, 4],  # features 1, node 1 in hyperedge 0
        #                                  [5, 6],  # features 2, node 1 in hyperedge 1
        #                                  [7, 8],  # features 3, node 2 in hyperedge 1
        #                                  [9, 10]] # features 4, node 3 in hyperedge 1
        #          -> incidence_aggregator pools features (0, 1) for hyperedge 0 and features (2, 3, 4) for hyperedge 1
        #          if aggregation == "maxmin":
        #          -> hyperedge_embeddings = [[max(1, 3) - min(1, 3), max(2, 4) - min(2, 4)],                # hyperedge 0
        #                                     [max(5, 7, 9) - min(5, 7, 9), max(6, 8, 10) - min(6, 8, 10)]]  # hyperedge 1
        #                                    shape: (num_hyperedges, hidden_channels)
        #         if aggregation == "mean":
        #         -> hyperedge_embeddings = [[mean(1, 3), mean(2, 4)],         # hyperedge 0
        #                                    [mean(5, 7, 9), mean(6, 8, 10)]]  # hyperedge 1
        #                                   shape: (num_hyperedges, hidden_channels)
        incidence_aggregator = HyperedgeAggregator(
            hyperedge_index=incidence_hyperedge_index,
            node_embeddings=incidence_embeddings,
        )

        match self.aggregation:
            case "maxmin":
                max_embeddings = incidence_aggregator.pool("max")
                min_embeddings = incidence_aggregator.pool("min")
                hyperedge_embeddings = max_embeddings - min_embeddings
            case _:
                hyperedge_embeddings = incidence_aggregator.pool("mean")

        # Decode: linear projection to scalar score per hyperedge
        # shape: (num_hyperedges, 1) -> squeeze -> (num_hyperedges,)
        return self.hyperedge_score(hyperedge_embeddings).squeeze(-1)