NN¶

class HyperedgeAggregator:
    """
    Pool node embeddings into hyperedge embeddings using the incidence structure.

    Each node-hyperedge incidence selects one node embedding row, then reduces
    those rows per hyperedge with the requested scatter aggregation.

    Args:
        hyperedge_index: Hyperedge incidence in COO format of size ``(2, num_incidences)``.
        node_embeddings: Node embedding matrix of size ``(num_nodes, num_channels)``.
        num_hyperedges: Optional explicit hyperedge count.
            When provided, the pooled output preserves empty hyperedges that do not appear in ``hyperedge_index``.
    """

    def __init__(
        self,
        hyperedge_index: Tensor,
        node_embeddings: Tensor,
        num_hyperedges: int | None = None,
    ):
        self.hyperedge_index_wrapper = HyperedgeIndex(hyperedge_index)
        self.node_embeddings = node_embeddings
        self.num_hyperedges = num_hyperedges

    def pool(self, aggregation: Literal["maxmin", "max", "min", "mean", "mul", "sum"]) -> Tensor:
        """
        Aggregate node embeddings for each hyperedge.

        ``hyperedge_index`` is the COO encoding of the nonzero entries of ``H``,
        so ``hyperedge_index[0, k] = v`` and ``hyperedge_index[1, k] = e`` means ``H[v, e] = 1`` for incidence ``k``.

        Let ``H`` be the binary incidence matrix of shape ``(num_nodes, num_hyperedges)``
        and let ``X`` be the node embedding matrix of shape ``(num_nodes, num_channels)``.
        This method pools node features into hyperedge features using the incidence pattern in ``H``:
        - ``aggregation="sum"`` computes the equivalent of the standard sparse matrix product ``H^T X``.
        - ``aggregation="mean"`` computes ``D_e^{-1} H^T X``, where ``D_e[e, e] = sum_v H[v, e]`` is the hyperedge cardinality matrix.
        - ``aggregation in {"max", "min", "mul"}`` uses the same sparsity pattern as ``H^T X``,
          but replaces the summation over incident nodes with a channel-wise ``max``, ``min``, or product reduction.
        - ``aggregation="maxmin"`` computes the channel-wise range ``max - min`` for each hyperedge.

        Examples:
            >>> hyperedge_index = [[0, 1, 2, 2, 3],
            ...                    [0, 0, 0, 1, 1]]
            >>> node_embeddings = [[1, 10], [2, 20], [3, 30], [4, 40]]
            >>> HyperedgeAggregator(hyperedge_index, node_embeddings).pool("mean")
            ... [[2, 20], [3.5, 35]]
            >>> HyperedgeAggregator(hyperedge_index, node_embeddings).pool("sum")
            ... [[6, 60], [7, 70]]
            >>> HyperedgeAggregator(hyperedge_index, node_embeddings).pool("max")
            ... [[3, 30], [4, 40]]
            >>> HyperedgeAggregator(hyperedge_index, node_embeddings).pool("maxmin")
            ... [[2, 20], [1, 10]]

        Args:
            aggregation: Reduction applied across the nodes belonging to each hyperedge.

        Returns:
            A hyperedge embedding matrix of shape ``(num_hyperedges, num_channels)``.
        """
        # Gather the embeddings for each incidence.
        # A node appearing in multiple hyperedges is repeated, once per incidence.
        # Example: node_embeddings = [[1, 10],  # node 0
        #                             [2, 20],  # node 1
        #                             [3, 30],  # node 2
        #                             [4, 40]]  # node 3
        #          -> all_node_ids = [0, 1, 2, 2, 3]
        #          -> incidence_node_embeddings = [[1, 10],  # node 0 for hyperedge 0
        #                                          [2, 20],  # node 1 for hyperedge 0
        #                                          [3, 30],  # node 2 for hyperedge 0
        #                                          [3, 30],  # node 2 for hyperedge 1
        #                                          [4, 40]]  # node 3 for hyperedge 1
        #             shape: (num_incidences, num_channels)
        incidence_node_embeddings = self.node_embeddings[self.hyperedge_index_wrapper.all_node_ids]

        # Scatter-aggregate node embeddings into hyperedge embeddings.
        # Example: with aggregation="sum":
        #          [[1+2+3, 10+20+30],  # hyperedge 0 contains node 0, 1, 2
        #          [3+4, 30+40]]        # hyperedge 1 contains node 2, 3
        #          shape: (num_hyperedges, num_channels)
        #          with aggregation="max":
        #          [[max(1, 2, 3), max(10, 20, 30)],  # hyperedge 0 contains node 0, 1, 2
        #           [max(3, 4), max(30, 40)]]         # hyperedge 1 contains node 2, 3
        #          shape: (num_hyperedges, num_channels)
        num_hyperedges = (
            self.num_hyperedges
            if self.num_hyperedges is not None
            else self.hyperedge_index_wrapper.num_hyperedges
        )

        if aggregation == "maxmin":
            return maxmin_scatter(
                src=incidence_node_embeddings,
                index=self.hyperedge_index_wrapper.all_hyperedge_ids,
                dim=0,  # scatter along the hyperedge dimension
                dim_size=num_hyperedges,
            )

        return scatter(
            src=incidence_node_embeddings,
            index=self.hyperedge_index_wrapper.all_hyperedge_ids,
            dim=0,  # scatter along the hyperedge dimension
            dim_size=num_hyperedges,
            reduce=aggregation,
        )

class NodeAggregator:
    """
    Pool hyperedge embeddings into node embeddings using the incidence structure.

    Each node-hyperedge incidence selects one hyperedge embedding row, then
    reduces those rows per node with the requested scatter aggregation.

    Args:
        hyperedge_index: Hyperedge incidence in COO format of size ``(2, num_incidences)``.
        hyperedge_embeddings: Hyperedge embedding matrix of size ``(num_hyperedges, num_channels)``.
        num_nodes: Optional explicit node count. When provided, the pooled output preserves isolated nodes that do not appear in ``hyperedge_index``.
    """

    def __init__(
        self,
        hyperedge_index: Tensor,
        hyperedge_embeddings: Tensor,
        num_nodes: int | None = None,
    ):
        self.hyperedge_index_wrapper = HyperedgeIndex(hyperedge_index)
        self.hyperedge_embeddings = hyperedge_embeddings
        self.num_nodes = num_nodes

    def pool(self, aggregation: Literal["maxmin", "max", "min", "mean", "mul", "sum"]) -> Tensor:
        """
        Aggregate hyperedge embeddings for each node.

        ``hyperedge_index`` is the COO encoding of the nonzero entries of ``H``,
        so ``hyperedge_index[0, k] = v`` and ``hyperedge_index[1, k] = e`` means ``H[v, e] = 1`` for incidence ``k``.

        Let ``H`` be the incidence matrix of shape ``(num_nodes, num_hyperedges)``
        and let ``E`` be the hyperedge embedding matrix of shape ``(num_hyperedges, num_channels)``.
        This method pools hyperedge features into node features using the incidence pattern in ``H``:
        - ``aggregation="sum"`` computes the equivalent of the standard sparse matrix product ``H E``.
        - ``aggregation="mean"`` computes ``D_v^{-1} H E``, where ``D_v[v, v] = sum_e H[v, e]`` is the node degree matrix.
        - ``aggregation in {"max", "min", "mul"}`` uses the same sparsity pattern as ``H E``,
          but replaces the summation over incident hyperedges with a channel-wise ``max``, ``min``, or product reduction.

        Examples:
            >>> hyperedge_index = [[0, 1, 1, 2],
            ...                    [0, 0, 1, 1]]
            >>> hyperedge_embeddings = [[10, 100], [20, 200]]
            >>> NodeAggregator(hyperedge_index, hyperedge_embeddings).pool("mean")
            ... [[10, 100], [15, 150], [20, 200]]
            >>> NodeAggregator(hyperedge_index, hyperedge_embeddings).pool("sum")
            ... [[10, 100], [30, 300], [20, 200]]
            >>> NodeAggregator(hyperedge_index, hyperedge_embeddings).pool("max")
            ... [[10, 100], [20, 200], [20, 200]]

        Args:
            aggregation: Reduction applied across the hyperedges incident to each node.

        Returns:
            A node embedding matrix of shape ``(num_nodes, num_channels)``.
        """
        # Gather the embeddings for each incidence.
        # A hyperedge appearing in multiple node incidences is repeated, once per incidence.
        # Example: hyperedge_embeddings = [[10, 100],  # hyperedge 0
        #                                  [20, 200]]  # hyperedge 1
        #          -> all_hyperedge_ids = [0, 0, 1, 1]
        #          -> incidence_hyperedge_embeddings = [[10, 100],   # hyperedge 0 for node 0
        #                                               [10, 100],   # hyperedge 0 for node 1
        #                                               [20, 200],   # hyperedge 1 for node 1
        #                                               [20, 200]]   # hyperedge 1 for node 2
        #             shape: (num_incidences, num_channels)
        incidence_hyperedge_embeddings = self.hyperedge_embeddings[
            self.hyperedge_index_wrapper.all_hyperedge_ids
        ]
        num_nodes = (
            self.num_nodes if self.num_nodes is not None else self.hyperedge_index_wrapper.num_nodes
        )

        if aggregation == "maxmin":
            return maxmin_scatter(
                src=incidence_hyperedge_embeddings,
                index=self.hyperedge_index_wrapper.all_node_ids,
                dim=0,  # scatter along the node dimension
                dim_size=num_nodes,
            )

        # Scatter-aggregate hyperedge embeddings into node embeddings.
        # Example: with aggregation="sum":
        #          [[10, 100],         # node 0 belongs to hyperedge 0
        #           [10+20, 100+200],  # node 1 belongs to hyperedge 0 and 1
        #           [20, 200]]         # node 2 belongs to hyperedge 1
        #          shape: (num_nodes, num_channels)
        #          with aggregation="max":
        #          [[10, 100],                     # node 0 belongs to hyperedge 0
        #           [max(10, 20), max(100, 200)],  # node 1 belongs to hyperedge 0 and 1
        #           [20, 200]]                     # node 2 belongs to hyperedge 1
        #         shape: (num_nodes, num_channels)
        return scatter(
            src=incidence_hyperedge_embeddings,
            index=self.hyperedge_index_wrapper.all_node_ids,
            dim=0,  # scatter along the node dimension
            dim_size=num_nodes,
            reduce=aggregation,
        )

class HGNNConv(nn.Module):
    """
    The HGNNConv layer 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/nn/convs/hypergraphs/hgnn_conv.html#HGNNConv>`_.

    Each layer performs: ``X' = sigma(L_HGNN X Theta)`` where ``L_HGNN = D_n^{-1/2} H D_e^{-1} H^T D_n^{-1/2}``
    is the hypergraph Laplacian computed from the incidence matrix H. This smooths node features through
    the hypergraph structure (nodes -> hyperedges -> nodes) without reducing to a pairwise graph.

    Unlike ``HyperGCNConv``, which uses a GCN Laplacian on a graph reduced from the hypergraph,
    ``HGNNConv`` operates entirely in hypergraph space and preserves all higher-order relationships.

    Args:
        in_channels: The number of input channels.
        out_channels: 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``, the layer will use batch normalization. Defaults to ``False``.
        drop_rate: If set to a positive number, the layer will use dropout. Defaults to ``0.5``.
        is_last: If set to ``True``, the layer will not apply the final activation and dropout functions. Defaults to ``False``.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        bias: bool = True,
        use_batch_normalization: bool = False,
        drop_rate: float = 0.5,
        is_last: bool = False,
    ):
        super().__init__()
        self.is_last = is_last
        self.batch_norm_1d = nn.BatchNorm1d(out_channels) if use_batch_normalization else None
        self.activation_fn = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(drop_rate)
        self.theta = nn.Linear(in_channels, out_channels, bias=bias)

    def forward(self, x: Tensor, hyperedge_index: Tensor) -> Tensor:
        """
        Apply one HGNN convolution layer: project features, smooth via hypergraph Laplacian,
        then apply activation, batch norm, and dropout (unless this is the last layer).

        The full per-layer formula is:
            ``X' = sigma( D_n^{-1/2} H D_e^{-1} H^T D_n^{-1/2} (X Theta) )``

        where the Laplacian ``L = D_n^{-1/2} H D_e^{-1} H^T D_n^{-1/2}`` is computed from
        the hyperedge_index and can be passed in precomputed as ``hgnn_laplacian_matrix``
        for efficiency when the hypergraph structure does not change across forward passes.

        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, out_channels)``.
        """
        x = self.theta(x)

        smoothing_matrix = HyperedgeIndex(hyperedge_index).get_sparse_hgnn_smoothing_matrix(
            num_nodes=x.size(0),
        )
        x = Hypergraph.smoothing_with_matrix(x, smoothing_matrix)

        if not self.is_last:
            x = self.activation_fn(x)
            if self.batch_norm_1d is not None:
                x = self.batch_norm_1d(x)
            x = self.dropout(x)

        return x

class HGNNPConv(nn.Module):
    """
    The HGNNPConv layer 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/nn/convs/hypergraphs/hgnnp_conv.html#HGNNPConv>`_.

    Each layer performs: ``X' = sigma(M_HGNN+ X Theta)`` where
    ``M_HGNN+ = D_v^{-1} H D_e^{-1} H^T`` is the HGNN+ smoothing matrix.

    Unlike ``HGNNConv``, which uses symmetric ``D_v^{-1/2}`` normalization for a
    spectral Laplacian, ``HGNNPConv`` uses plain inverse degrees and performs
    two-stage mean aggregation: nodes -> hyperedges -> nodes.

    Args:
        in_channels: The number of input channels.
        out_channels: 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``, the layer will use batch normalization. Defaults to ``False``.
        drop_rate: If set to a positive number, the layer will use dropout. Defaults to ``0.5``.
        is_last: If set to ``True``, the layer will not apply the final activation and dropout functions. Defaults to ``False``.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        bias: bool = True,
        use_batch_normalization: bool = False,
        drop_rate: float = 0.5,
        is_last: bool = False,
    ):
        super().__init__()
        self.is_last = is_last
        self.batch_norm_1d = nn.BatchNorm1d(out_channels) if use_batch_normalization else None
        self.activation_fn = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(drop_rate)
        self.theta = nn.Linear(in_channels, out_channels, bias=bias)

    def forward(self, x: Tensor, hyperedge_index: Tensor) -> Tensor:
        """
        Apply one HGNN+ convolution layer using row-stochastic hypergraph smoothing.

        The full per-layer formula is:
            ``X' = sigma( D_v^{-1} H D_e^{-1} H^T (X Theta) )``

        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, out_channels)``.
        """
        x = self.theta(x)

        smoothing_matrix = HyperedgeIndex(hyperedge_index).get_sparse_hgnnp_smoothing_matrix(
            num_nodes=x.size(0),
        )
        x = Hypergraph.smoothing_with_matrix(x, smoothing_matrix)

        if not self.is_last:
            x = self.activation_fn(x)
            if self.batch_norm_1d is not None:
                x = self.batch_norm_1d(x)
            x = self.dropout(x)

        return x

class HNHNConv(nn.Module):
    """
    The HNHNConv layer 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/nn/convs/hypergraphs/hnhn_conv.html#HNHNConv>`_.

    Args:
        in_channels: The number of input channels.
        out_channels: 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``, the layer will use batch normalization. Defaults to ``False``.
        drop_rate: If set to a positive number, the layer will use dropout. Defaults to ``0.5``.
        is_last: If set to ``True``, the layer will not apply the final activation and dropout functions. Defaults to ``False``.
    """

    __AGGREGATION: Literal["mean"] = "mean"

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        bias: bool = True,
        use_batch_normalization: bool = False,
        drop_rate: float = 0.5,
        is_last: bool = False,
    ):
        super().__init__()
        self.is_last = is_last
        self.batch_norm_1d = nn.BatchNorm1d(out_channels) if use_batch_normalization else None
        self.activation_fn = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(drop_rate)
        self.theta_v2e = nn.Linear(in_channels, out_channels, bias=bias)
        self.theta_e2v = nn.Linear(out_channels, out_channels, bias=bias)

    def forward(self, x: Tensor, hyperedge_index: Tensor) -> Tensor:
        """
        Apply one HNHN convolution layer using two learned projections around
        node-to-hyperedge and hyperedge-to-node mean aggregation.

        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, out_channels)``.
        """
        x = self.theta_v2e(x)

        hyperedge_embeddings = HyperedgeAggregator(hyperedge_index, x).pool(self.__AGGREGATION)
        hyperedge_embeddings = self.activation_fn(hyperedge_embeddings)
        hyperedge_embeddings = self.theta_e2v(hyperedge_embeddings)

        x = NodeAggregator(
            hyperedge_index=hyperedge_index,
            hyperedge_embeddings=hyperedge_embeddings,
            num_nodes=x.size(0),
        ).pool(self.__AGGREGATION)

        if not self.is_last:
            x = self.activation_fn(x)
            if self.batch_norm_1d is not None:
                x = self.batch_norm_1d(x)
            x = self.dropout(x)

        return x

class HyperGCNConv(nn.Module):
    """
    The HyperGCNConv layer 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).
    Reference implementation: `source <https://deephypergraph.readthedocs.io/en/latest/_modules/dhg/nn/convs/hypergraphs/hypergcn_conv.html#HyperGCNConv>`_.

    Args:
        in_channels: The number of input channels.
        out_channels: 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``, the layer will use batch normalization. Defaults to ``False``.
        drop_rate: If set to a positive number, the layer will use dropout. Defaults to ``0.5``.
        use_mediator: Whether to use mediator to transform the hyperedges to edges in the graph. Defaults to ``False``.
        is_last: If set to ``True``, the layer will not apply the final activation and dropout functions. Defaults to ``False``.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        bias: bool = True,
        use_batch_normalization: bool = False,
        drop_rate: float = 0.5,
        use_mediator: bool = False,
        is_last: bool = False,
    ):
        super().__init__()
        self.is_last = is_last
        self.use_mediator = use_mediator
        self.batch_norm_1d = nn.BatchNorm1d(out_channels) if use_batch_normalization else None
        self.activation_fn = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(drop_rate)

        # θ is the learnable weight matrix (as in the HyperGCN paper),
        # it projects node features from in_channels to out_channels and learns how to mix feature channels
        self.theta = nn.Linear(in_channels, out_channels, bias=bias)

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

        Args:
            x: Input node feature matrix. Size ``(num_nodes, in_channels)``.
            hyperedge_index: Hyperedge indices representing the hypergraph structure. Size ``(2, num_hyperedges)``.
            gcn_laplacian_matrix: Optional precomputed normalized GCN Laplacian matrix. Size ``(num_nodes, num_nodes)``. Defaults to ``None``.
                If provided, it will be used directly for smoothing, so we can skip computing it from edge_index.

        Returns:
            The output node feature matrix. Size ``(num_nodes, out_channels)``.
        """
        x = self.theta(x)

        if gcn_laplacian_matrix is not None:
            x = Graph.smoothing_with_laplacian_matrix(x, gcn_laplacian_matrix)
        else:
            edge_index, edge_weights = HyperedgeIndex(
                hyperedge_index
            ).reduce_to_edge_index_on_random_direction(
                x=x,
                with_mediators=self.use_mediator,
                return_weights=True,
            )

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

            x = Graph.smoothing_with_laplacian_matrix(x, normalized_gcn_laplacian_matrix)

        if not self.is_last:
            x = self.activation_fn(x)
            if self.batch_norm_1d is not None:
                x = self.batch_norm_1d(x)
            x = self.dropout(x)

        return x

class NodeEnricher(Enricher, ABC):
    """
    Base class for node enrichers.
    """

    pass

class HyperedgeEnricher(Enricher, ABC):
    """
    Base class for hyperedge enrichers.
    """

    pass

class FillValueHyperedgeAttrsEnricher(HyperedgeAttrsEnricher):
    """
    Generates simple hyperedge attributes by filling them with a constant value.

    Args:
        cache_dir: Directory for saving/loading cached features. If ``None``, caching is disabled.
        fill_value: The constant value to fill the hyperedge attributes with. Defaults to ``1.0``.
    """

    def __init__(
        self,
        cache_dir: str | None = None,
        fill_value: float = 1.0,
    ):
        super().__init__(cache_dir=cache_dir)
        self.fill_value = fill_value

    def enrich(self, hyperedge_index: Tensor) -> Tensor:
        """
        Generate hyperedge attributes.

        Args:
            hyperedge_index: Hyperedge index tensor of shape ``(2, num_hyperedges)``.

        Returns:
            Tensor of shape ``(num_hyperedges, 1)`` containing the generated attribute for each hyperedge.
        """
        num_hyperedges = HyperedgeIndex(hyperedge_index).num_hyperedges
        hyperedge_attrs = torch.full(
            size=(num_hyperedges, 1),
            fill_value=self.fill_value,
            device=hyperedge_index.device,
        )
        return hyperedge_attrs

class ABHyperedgeWeightsEnricher(HyperedgeWeightsEnricher):
    """
    Generates hyperedge weights based on the number of nodes in each hyperedge.

    Args:
        cache_dir: Directory for saving/loading cached features. If ``None``, caching is disabled.
        alpha: Scaling factor for the random component added to weights. Must be between 0.0 and 1.0.
        beta: If provided, the random component is alpha * beta. If None, no random component is added.
    """

    def __init__(
        self,
        cache_dir: str | None = None,
        alpha: float = 1.0,
        beta: float | None = None,
    ):
        super().__init__(cache_dir=cache_dir)
        if alpha < 0.0 or alpha > 1.0:
            raise ValueError("Alpha must be between 0.0 and 1.0.")

        self.alpha = alpha
        self.beta = beta

    def enrich(self, hyperedge_index: Tensor) -> Tensor:
        """
        Compute edge weights as the number of nodes in each hyperedge.

        Args:
            hyperedge_index: Hyperedge index tensor of shape ``(2, num_hyperedges)``.

        Returns:
            Tensor of shape ``(num_hyperedges,)`` containing the weight of each hyperedge.
        """
        # Count the number of nodes in each hyperedge by counting occurrences of each hyperedge index.
        # Example: if hyperedge_index[1] = [0, 0, 1, 1, 1], then we have 2 nodes in hyperedge 0 and 3 nodes in hyperedge 1.
        num_hyperedges = int(hyperedge_index[1].max().item()) + 1
        weights = torch.bincount(hyperedge_index[1], minlength=num_hyperedges).float()

        random_alpha = random.uniform(0, self.alpha)
        if self.beta is not None:
            weights += random_alpha * self.beta
        return weights