HLP¶

class CommonNeighborsHlpModule(HlpModule):
    """
    A LightningModule for the CommonNeighbors model with optional negative sampling.

    Args:
        aggregation: The aggregation method for common neighbors ("mean", "min", or "sum").
        decoder: An optional decoder module. Defaults to :class:`CommonNeighbors`.
        loss_fn: An optional loss function. Defaults to ``BCEWithLogitsLoss``.
        metrics: An optional dictionary of metric functions.
    """

    def __init__(
        self,
        train_hyperedge_index: Tensor,
        aggregation: Literal["mean", "min", "sum"] = "mean",
        decoder: nn.Module | None = None,
        loss_fn: nn.Module | None = None,
        metrics: MetricCollection | None = None,
    ):
        super().__init__(
            decoder=decoder if decoder is not None else CommonNeighbors(aggregation),
            loss_fn=loss_fn if loss_fn is not None else nn.BCEWithLogitsLoss(),
            metrics=metrics,
        )

        # Pre-compute neighbors of training nodes based on training edges only
        # to create a "known world" for the model to make predictions from
        self.node_to_neighbors = Hypergraph.from_hyperedge_index(
            train_hyperedge_index
        ).neighbors_of_all()

        # Disable automatic optimization since there is no training
        self.automatic_optimization = False

    def forward(self, hyperedge_index: Tensor) -> Tensor:
        """
        Compute common neighbor scores for the given hyperedges.

        Args:
            hyperedge_index: Tensor containing incidence information for the hyperedges to score.
        """
        return self.decoder(hyperedge_index, self.node_to_neighbors)

    def on_fit_start(self) -> None:
        """Warn users if they are running unnecessary training epochs."""
        if self.trainer.max_epochs is None or self.trainer.max_epochs > 0:
            warnings.warn(
                f"{self.__class__.__name__} is a non-trainable heuristic model. "
                "No optimization occurs. Set max_epochs=0 in your trainer for instant evaluation.",
                UserWarning,
                stacklevel=2,
            )

    def training_step(self, batch: HData, batch_idx: int) -> Tensor:
        return torch.tensor(0.0, device=self.device)

    def validation_step(self, batch: HData, batch_idx: int) -> Tensor:
        return torch.tensor(0.0, device=self.device)

    def test_step(self, batch: HData, batch_idx: int) -> Tensor:
        return self.__step(batch, stage=Stage.TEST)

    def predict_step(self, batch: HData, batch_idx: int) -> Tensor:
        return self.forward(batch.hyperedge_index)

    def configure_optimizers(self):
        # No training, so no optimizers needed
        return None

    def __step(self, batch: HData, stage: Stage) -> Tensor:
        """
        Shared evaluation logic for all stages.

        Args:
            batch: :class:`HData` object containing the hypergraph.
            stage: The current stage of evaluation (e.g., ``Stage.TRAIN``, ``Stage.VAL``, ``Stage.TEST``).

        Returns:
            The computed loss.
        """
        scores = self.forward(batch.hyperedge_index)
        labels = batch.y

        # We need to use the number of hyperedges as batch size for logging purposes,
        # since each hyperedge is a separate prediction
        batch_size = batch.num_hyperedges

        loss = self._compute_loss(scores, labels, batch_size, stage)
        self._compute_metrics(scores, labels, batch_size, stage)

        return loss

class GCNEncoderConfig(TypedDict):
    """
    Configuration for the GCN encoder in GCNHlpModule.

    Args:
        in_channels: Number of input features per node.
        out_channels: Number of output features (embedding size) per node.
        hidden_channels: Number of hidden units in the intermediate GCN layers.
        num_layers: Number of GCN layers. Defaults to ``2``.
        drop_rate: Dropout rate applied after each hidden GCN layer. Defaults to ``0.0``.
        bias: Whether to include bias terms. Defaults to ``True``.
        improved: Whether to use the improved GCN normalization. Defaults to ``False``.
        add_self_loops: Whether to add self-loops before convolution. Defaults to ``True``.
        normalize: Whether to normalize the adjacency matrix in ``GCNConv``. Defaults to ``True``.
        cached: Whether to cache the normalized graph in ``GCNConv``. Defaults to ``False``.
        graph_reduction_strategy: Strategy for reducing the hypergraph to a graph. Defaults to ``"clique_expansion"``.
        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.
    """

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

class GCNHlpModule(HlpModule):
    """
    A LightningModule for GCN-based HLP.

    Uses a graph reduction of the input hypergraph to run GCN over nodes,
    aggregates node embeddings per hyperedge, and scores each hyperedge with a linear decoder.

    Args:
        encoder_config: Configuration for the GCN encoder.
        aggregation: Method to aggregate node embeddings per hyperedge. Defaults to ``"mean"``.
        loss_fn: Loss function. Defaults to ``BCEWithLogitsLoss``.
        lr: Learning rate for the optimizer. Defaults to ``0.001``.
        weight_decay: L2 regularization. Defaults to ``0.0``.
        metrics: Optional metric collection for evaluation.
    """

    def __init__(
        self,
        encoder_config: GCNEncoderConfig,
        aggregation: Literal["mean", "max", "min", "sum"] = "mean",
        loss_fn: nn.Module | None = None,
        lr: float = 0.001,
        weight_decay: float = 0.0,
        metrics: MetricCollection | None = None,
    ):
        encoder = GCN(
            in_channels=encoder_config["in_channels"],
            out_channels=encoder_config["out_channels"],
            hidden_channels=encoder_config.get("hidden_channels"),
            num_layers=encoder_config.get("num_layers", 2),
            drop_rate=encoder_config.get("drop_rate", 0.0),
            bias=encoder_config.get("bias", True),
            activation_fn=encoder_config.get("activation_fn"),
            activation_fn_kwargs=encoder_config.get("activation_fn_kwargs"),
            improved=encoder_config.get("improved", False),
            add_self_loops=encoder_config.get("add_self_loops", True),
            normalize=encoder_config.get("normalize", True),
            cached=encoder_config.get("cached", False),
        )
        decoder = SLP(in_channels=encoder_config["out_channels"], out_channels=1)

        super().__init__(
            encoder=encoder,
            decoder=decoder,
            loss_fn=loss_fn if loss_fn is not None else nn.BCEWithLogitsLoss(),
            metrics=metrics,
        )

        self.encoder_config = encoder_config
        self.aggregation = aggregation
        self.lr = lr
        self.weight_decay = weight_decay

    def forward(self, x: Tensor, hyperedge_index: Tensor) -> Tensor:
        """
        Reduce the hypergraph to a graph, encode nodes with GCN, aggregate per hyperedge, and score.

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

        Returns:
            Logit scores of shape ``(num_hyperedges,)``.
        """
        if self.encoder is None:
            raise ValueError("Encoder is not defined for this HLP module.")

        # Reduce hypergraph to graph and remove self-loops
        reduced_edge_index = HyperedgeIndex(hyperedge_index).reduce(
            strategy=self.encoder_config.get("graph_reduction_strategy", "clique_expansion")
        )
        edge_index = EdgeIndex(reduced_edge_index).remove_selfloops().item

        # Encode nodes with GCN
        node_embeddings: Tensor = self.encoder(x, edge_index)

        # Aggregate node embeddings per hyperedge
        hyperedge_embeddings = HyperedgeAggregator(hyperedge_index, node_embeddings).pool(
            self.aggregation
        )

        return self.decoder(hyperedge_embeddings).squeeze(-1)

    def training_step(self, batch: HData, batch_idx: int) -> Tensor:
        return self.__eval_step(batch, Stage.TRAIN)

    def validation_step(self, batch: HData, batch_idx: int) -> Tensor:
        return self.__eval_step(batch, Stage.VAL)

    def test_step(self, batch: HData, batch_idx: int) -> Tensor:
        return self.__eval_step(batch, Stage.TEST)

    def predict_step(self, batch: HData, batch_idx: int) -> Tensor:
        return self.forward(batch.x, batch.hyperedge_index)

    def configure_optimizers(self):
        return optim.Adam(self.parameters(), lr=self.lr, weight_decay=self.weight_decay)

    def __eval_step(self, batch: HData, stage: Stage) -> Tensor:
        scores = self.forward(batch.x, batch.hyperedge_index)
        labels = batch.y
        batch_size = batch.num_hyperedges

        loss = self._compute_loss(scores, labels, batch_size, stage)
        self._compute_metrics(scores, labels, batch_size, stage)
        return loss

class HGNNHlpModule(HlpModule):
    """
    A LightningModule for HGNN-based Hyperedge Link Prediction.

    Uses HGNN as an encoder to produce structure-aware node embeddings via
    spectral hypergraph convolution, aggregates them per hyperedge,
    and scores each hyperedge with a linear decoder.

    Args:
        encoder_config: Configuration for the HGNN encoder.
        aggregation: Method to aggregate node embeddings per hyperedge. Defaults to ``"mean"``.
        loss_fn: Loss function. Defaults to ``BCEWithLogitsLoss``.
        lr: Learning rate for the optimizer. Defaults to ``0.01``.
        weight_decay: L2 regularization. Defaults to ``5e-4``.
        metrics: Optional metric collection for evaluation.
    """

    def __init__(
        self,
        encoder_config: HGNNEncoderConfig,
        aggregation: Literal["mean", "max", "min", "sum"] = "mean",
        loss_fn: nn.Module | None = None,
        lr: float = 0.001,
        weight_decay: float = 5e-4,
        metrics: MetricCollection | None = None,
    ):
        encoder = HGNN(
            in_channels=encoder_config["in_channels"],
            hidden_channels=encoder_config["hidden_channels"],
            num_classes=encoder_config["out_channels"],
            bias=encoder_config.get("bias", True),
            use_batch_normalization=encoder_config.get("use_batch_normalization", False),
            drop_rate=encoder_config.get("drop_rate", 0.5),
        )
        decoder = SLP(in_channels=encoder_config["out_channels"], out_channels=1)

        super().__init__(
            encoder=encoder,
            decoder=decoder,
            loss_fn=loss_fn if loss_fn is not None else nn.BCEWithLogitsLoss(),
            metrics=metrics,
        )

        self.aggregation = aggregation
        self.lr = lr
        self.weight_decay = weight_decay

    def forward(self, x: Tensor, hyperedge_index: Tensor) -> Tensor:
        """
        Run the full HGNN-based hyperedge link prediction pipeline.

        The pipeline has three stages:
        1. Encode: HGNN applies two rounds of ``D_n^{-1/2} H D_e^{-1} H^T D_n^{-1/2}``
           smoothing to propagate information through the hypergraph topology (nodes ->
           hyperedges -> nodes). The output is a structure-aware node embedding matrix of
           shape ``(num_nodes, out_channels)``.
        2. Aggregate: For each hyperedge being scored, pool the embeddings of its member
           nodes using the configured strategy (mean/max/min/sum). This produces a hyperedge
           embedding that summarizes the collective representation of the hyperedge's nodes.
           Shape: ``(num_hyperedges, out_channels)``.
        3. Decode: A single linear layer (SLP) projects each hyperedge embedding to a
           scalar score representing the likelihood that the hyperedge is a real (positive)
           hyperedge. Shape: ``(num_hyperedges,)``.

        Examples:
            Given 5 nodes with 8 features and 2 hyperedges::

                >>> x.shape  # (5, 8) - all nodes in the hypergraph
                >>> hyperedge_index = [[0, 1, 2, 3, 4],  # node IDs
                ...                    [0, 0, 0, 1, 1]]  # hyperedge IDs

            The forward pass:
                1. HGNN encodes all 5 nodes using the hypergraph Laplacian.
                   ``node_embeddings.shape = (5, out_channels)``
                2. Aggregate per hyperedge:
                   - hyperedge 0: pool(emb[0], emb[1], emb[2])
                   - hyperedge 1: pool(emb[3], emb[4])
                   ``hyperedge_embeddings.shape = (2, out_channels)``
                3. Decode: one scalar per hyperedge -> ``scores.shape = (2,)``

        Args:
            x: Node feature matrix of shape ``(num_nodes, in_channels)``.
                Must contain **all** nodes referenced in ``hyperedge_index``.
            hyperedge_index: Hyperedge connectivity of shape ``(2, num_incidences)``,
                with row 0 containing global node IDs and row 1 hyperedge IDs.

        Returns:
            Logit scores of shape ``(num_hyperedges,)``. Pass through sigmoid to get
            probabilities, or use directly with ``BCEWithLogitsLoss``.
        """
        if self.encoder is None:
            raise ValueError("Encoder is not defined for this HLP module.")

        # Encode: two-hop HGNN smoothing (nodes -> hyperedges -> nodes), no graph reduction
        # Example: x: (num_nodes, in_channels)
        #          -> node_embeddings: (num_nodes, out_channels)
        node_embeddings: Tensor = self.encoder(x, hyperedge_index)

        # Aggregate: pool node embeddings per hyperedge
        # shape: (num_hyperedges, out_channels)
        hyperedge_embeddings = HyperedgeAggregator(hyperedge_index, node_embeddings).pool(
            self.aggregation
        )

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

    def training_step(self, batch: HData, batch_idx: int) -> Tensor:
        return self.__eval_step(batch, Stage.TRAIN)

    def validation_step(self, batch: HData, batch_idx: int) -> Tensor:
        return self.__eval_step(batch, Stage.VAL)

    def test_step(self, batch: HData, batch_idx: int) -> Tensor:
        return self.__eval_step(batch, Stage.TEST)

    def predict_step(self, batch: HData, batch_idx: int) -> Tensor:
        return self.forward(batch.x, batch.hyperedge_index)

    def configure_optimizers(self):
        return optim.Adam(self.parameters(), lr=self.lr, weight_decay=self.weight_decay)

    def __eval_step(self, batch: HData, stage: Stage) -> Tensor:
        scores = self.forward(batch.x, batch.hyperedge_index)
        labels = batch.y
        batch_size = batch.num_hyperedges

        loss = self._compute_loss(scores, labels, batch_size, stage)
        self._compute_metrics(scores, labels, batch_size, stage)
        return loss

class HGNNEncoderConfig(TypedDict):
    """
    Configuration for the HGNN encoder in HGNNHlpModule.

    Args:
        in_channels: Number of input features per node.
        hidden_channels: Number of hidden units in the intermediate HGNN layer.
        out_channels: Number of output features (embedding size) per node.
        bias: Whether to include bias terms. Defaults to ``True``.
        use_batch_normalization: Whether to use batch normalization. Defaults to ``False``.
        drop_rate: Dropout rate. Defaults to ``0.5``.
    """

    in_channels: int
    hidden_channels: int
    out_channels: int
    bias: NotRequired[bool]
    use_batch_normalization: NotRequired[bool]
    drop_rate: NotRequired[float]

class HNHNEncoderConfig(TypedDict):
    """
    Configuration for the HNHN encoder in HNHNHlpModule.

    Args:
        in_channels: Number of input features per node.
        hidden_channels: Number of hidden units in the intermediate HNHN layer.
        out_channels: Number of output features (embedding size) per node.
        bias: Whether to include bias terms. Defaults to ``True``.
        use_batch_normalization: Whether to use batch normalization. Defaults to ``False``.
        drop_rate: Dropout rate. Defaults to ``0.5``.
    """

    in_channels: int
    hidden_channels: int
    out_channels: int
    bias: NotRequired[bool]
    use_batch_normalization: NotRequired[bool]
    drop_rate: NotRequired[float]

class HNHNHlpModule(HlpModule):
    """
    A LightningModule for HNHN-based Hyperedge Link Prediction.

    Uses HNHN as an encoder to produce node embeddings through explicit
    hyperedge neurons, aggregates them per hyperedge, and scores each
    hyperedge with a linear decoder.

    Args:
        encoder_config: Configuration for the HNHN encoder.
        aggregation: Method to aggregate node embeddings per hyperedge. Defaults to ``"mean"``.
        loss_fn: Loss function. Defaults to ``BCEWithLogitsLoss``.
        lr: Learning rate for the optimizer. Defaults to ``0.01``.
        weight_decay: L2 regularization. Defaults to ``5e-4``.
        scheduler_step_size: Step size for learning rate scheduler. Defaults to ``100``.
        scheduler_gamma: Multiplicative factor for learning rate decay. Defaults to ``0.51``.
        metrics: Optional metric collection for evaluation.
    """

    def __init__(
        self,
        encoder_config: HNHNEncoderConfig,
        aggregation: Literal["mean", "max", "min", "sum"] = "mean",
        loss_fn: nn.Module | None = None,
        lr: float = 0.01,
        weight_decay: float = 5e-4,
        scheduler_step_size: int = 100,
        scheduler_gamma: float = 0.51,
        metrics: MetricCollection | None = None,
    ):
        encoder = HNHN(
            in_channels=encoder_config["in_channels"],
            hidden_channels=encoder_config["hidden_channels"],
            num_classes=encoder_config["out_channels"],
            bias=encoder_config.get("bias", True),
            use_batch_normalization=encoder_config.get("use_batch_normalization", False),
            drop_rate=encoder_config.get("drop_rate", 0.5),
        )
        decoder = SLP(in_channels=encoder_config["out_channels"], out_channels=1)

        super().__init__(
            encoder=encoder,
            decoder=decoder,
            loss_fn=loss_fn if loss_fn is not None else nn.BCEWithLogitsLoss(),
            metrics=metrics,
        )

        self.aggregation = aggregation
        self.lr = lr
        self.weight_decay = weight_decay
        self.scheduler_step_size = scheduler_step_size
        self.scheduler_gamma = scheduler_gamma

    def forward(self, x: Tensor, hyperedge_index: Tensor) -> Tensor:
        """
        Run the full HNHN-based hyperedge link prediction pipeline.

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

        Returns:
            Logit scores of shape ``(num_hyperedges,)``.
        """
        if self.encoder is None:
            raise ValueError("Encoder is not defined for this HLP module.")

        node_embeddings: Tensor = self.encoder(x, hyperedge_index)
        hyperedge_embeddings = HyperedgeAggregator(hyperedge_index, node_embeddings).pool(
            self.aggregation
        )
        scores: Tensor = self.decoder(hyperedge_embeddings).squeeze(-1)
        return scores

    def training_step(self, batch: HData, batch_idx: int) -> Tensor:
        return self.__eval_step(batch, Stage.TRAIN)

    def validation_step(self, batch: HData, batch_idx: int) -> Tensor:
        return self.__eval_step(batch, Stage.VAL)

    def test_step(self, batch: HData, batch_idx: int) -> Tensor:
        return self.__eval_step(batch, Stage.TEST)

    def predict_step(self, batch: HData, batch_idx: int) -> Tensor:
        return self.forward(batch.x, batch.hyperedge_index)

    def configure_optimizers(self):
        optimizer = optim.Adam(self.parameters(), lr=self.lr, weight_decay=self.weight_decay)
        scheduler = optim.lr_scheduler.StepLR(
            optimizer, step_size=self.scheduler_step_size, gamma=self.scheduler_gamma
        )
        return [optimizer], [scheduler]

    def __eval_step(self, batch: HData, stage: Stage) -> Tensor:
        scores = self.forward(batch.x, batch.hyperedge_index)
        labels = batch.y
        batch_size = batch.num_hyperedges

        loss = self._compute_loss(scores, labels, batch_size, stage)
        self._compute_metrics(scores, labels, batch_size, stage)
        return loss

class HGNNPEncoderConfig(TypedDict):
    """
    Configuration for the HGNN+ encoder in HGNNPHlpModule.

    Args:
        in_channels: Number of input features per node.
        hidden_channels: Number of hidden units in the intermediate HGNN+ layer.
        out_channels: Number of output features (embedding size) per node.
        bias: Whether to include bias terms. Defaults to ``True``.
        use_batch_normalization: Whether to use batch normalization. Defaults to ``False``.
        drop_rate: Dropout rate. Defaults to ``0.5``.
    """

    in_channels: int
    hidden_channels: int
    out_channels: int
    bias: NotRequired[bool]
    use_batch_normalization: NotRequired[bool]
    drop_rate: NotRequired[float]

class HGNNPHlpModule(HlpModule):
    """
    A LightningModule for HGNN+-based Hyperedge Link Prediction.

    Uses HGNN+ as an encoder to produce structure-aware node embeddings via
    row-stochastic hypergraph convolution, aggregates them per hyperedge,
    and scores each hyperedge with a linear decoder.

    Args:
        encoder_config: Configuration for the HGNN+ encoder.
        aggregation: Method to aggregate node embeddings per hyperedge. Defaults to ``"mean"``.
        loss_fn: Loss function. Defaults to ``BCEWithLogitsLoss``.
        lr: Learning rate for the optimizer. Defaults to ``0.01``.
        weight_decay: L2 regularization. Defaults to ``5e-4``.
        metrics: Optional metric collection for evaluation.
    """

    def __init__(
        self,
        encoder_config: HGNNPEncoderConfig,
        aggregation: Literal["mean", "max", "min", "sum"] = "mean",
        loss_fn: nn.Module | None = None,
        lr: float = 0.01,
        weight_decay: float = 5e-4,
        metrics: MetricCollection | None = None,
    ):
        encoder = HGNNP(
            in_channels=encoder_config["in_channels"],
            hidden_channels=encoder_config["hidden_channels"],
            num_classes=encoder_config["out_channels"],
            bias=encoder_config.get("bias", True),
            use_batch_normalization=encoder_config.get("use_batch_normalization", False),
            drop_rate=encoder_config.get("drop_rate", 0.5),
        )
        decoder = SLP(in_channels=encoder_config["out_channels"], out_channels=1)

        super().__init__(
            encoder=encoder,
            decoder=decoder,
            loss_fn=loss_fn if loss_fn is not None else nn.BCEWithLogitsLoss(),
            metrics=metrics,
        )

        self.aggregation = aggregation
        self.lr = lr
        self.weight_decay = weight_decay

    def forward(self, x: Tensor, hyperedge_index: Tensor) -> Tensor:
        """
        Run the full HGNN+-based hyperedge link prediction pipeline.

        The pipeline has three stages:
        1. Encode: HGNN+ applies two rounds of ``D_v^{-1} H D_e^{-1} H^T``
           smoothing to propagate information through the hypergraph topology with
           two-stage mean aggregation. The output is a structure-aware node
           embedding matrix of shape ``(num_nodes, out_channels)``.
        2. Aggregate: For each hyperedge being scored, pool the embeddings of its member
           nodes using the configured strategy (mean/max/min/sum). This produces a hyperedge
           embedding of shape ``(num_hyperedges, out_channels)``.
        3. Decode: A single linear layer projects each hyperedge embedding to a
           scalar score. Shape: ``(num_hyperedges,)``.

        Args:
            x: Node feature matrix of shape ``(num_nodes, in_channels)``.
                Must contain **all** nodes referenced in ``hyperedge_index``.
            hyperedge_index: Hyperedge connectivity of shape ``(2, num_incidences)``,
                with row 0 containing global node IDs and row 1 hyperedge IDs.

        Returns:
            Logit scores of shape ``(num_hyperedges,)``.
        """
        if self.encoder is None:
            raise ValueError("Encoder is not defined for this HLP module.")

        # Encode: produce node embeddings using HGNN+, no graph reduction is applied
        # Example: x: (num_nodes, in_channels)
        #          -> node_embeddings: (num_nodes, out_channels), out_channels)
        node_embeddings: Tensor = self.encoder(x, hyperedge_index)

        # Aggregate: pool node embeddings per hyperedge
        # shape: (num_hyperedges, out_channels)
        hyperedge_embeddings = HyperedgeAggregator(hyperedge_index, node_embeddings).pool(
            self.aggregation
        )

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

    def training_step(self, batch: HData, batch_idx: int) -> Tensor:
        return self.__eval_step(batch, Stage.TRAIN)

    def validation_step(self, batch: HData, batch_idx: int) -> Tensor:
        return self.__eval_step(batch, Stage.VAL)

    def test_step(self, batch: HData, batch_idx: int) -> Tensor:
        return self.__eval_step(batch, Stage.TEST)

    def predict_step(self, batch: HData, batch_idx: int) -> Tensor:
        return self.forward(batch.x, batch.hyperedge_index)

    def configure_optimizers(self):
        return optim.Adam(self.parameters(), lr=self.lr, weight_decay=self.weight_decay)

    def __eval_step(self, batch: HData, stage: Stage) -> Tensor:
        scores = self.forward(batch.x, batch.hyperedge_index)
        labels = batch.y
        batch_size = batch.num_hyperedges

        loss = self._compute_loss(scores, labels, batch_size, stage)
        self._compute_metrics(scores, labels, batch_size, stage)
        return loss