pyg-nightly 2.6.0.dev20240704__py3-none-any.whl → 2.8.0.dev20251207__py3-none-any.whl

torch_geometric/utils/_scatter.py

@@ -4,174 +4,138 @@ import torch
 from torch import Tensor
 
 import torch_geometric.typing
-from torch_geometric import is_compiling, warnings
+from torch_geometric import is_compiling, is_in_onnx_export, warnings
 from torch_geometric.typing import torch_scatter
 from torch_geometric.utils.functions import cumsum
 
-if torch_geometric.typing.WITH_PT112:  # pragma: no cover
-
-    warnings.filterwarnings('ignore', '.*is in beta and the API may change.*')
-
-    def scatter(
-        src: Tensor,
-        index: Tensor,
-        dim: int = 0,
-        dim_size: Optional[int] = None,
-        reduce: str = 'sum',
-    ) -> Tensor:
-        r"""Reduces all values from the :obj:`src` tensor at the indices
-        specified in the :obj:`index` tensor along a given dimension
-        :obj:`dim`. See the `documentation
-        <https://pytorch-scatter.readthedocs.io/en/latest/functions/
-        scatter.html>`__ of the :obj:`torch_scatter` package for more
-        information.
-
-        Args:
-            src (torch.Tensor): The source tensor.
-            index (torch.Tensor): The index tensor.
-            dim (int, optional): The dimension along which to index.
-                (default: :obj:`0`)
-            dim_size (int, optional): The size of the output tensor at
-                dimension :obj:`dim`. If set to :obj:`None`, will create a
-                minimal-sized output tensor according to
-                :obj:`index.max() + 1`. (default: :obj:`None`)
-            reduce (str, optional): The reduce operation (:obj:`"sum"`,
-                :obj:`"mean"`, :obj:`"mul"`, :obj:`"min"` or :obj:`"max"`,
-                :obj:`"any"`). (default: :obj:`"sum"`)
-        """
-        if isinstance(index, Tensor) and index.dim() != 1:
-            raise ValueError(f"The `index` argument must be one-dimensional "
-                             f"(got {index.dim()} dimensions)")
-
-        dim = src.dim() + dim if dim < 0 else dim
-
-        if isinstance(src, Tensor) and (dim < 0 or dim >= src.dim()):
-            raise ValueError(f"The `dim` argument must lay between 0 and "
-                             f"{src.dim() - 1} (got {dim})")
-
-        if dim_size is None:
-            dim_size = int(index.max()) + 1 if index.numel() > 0 else 0
-
-        # For now, we maintain various different code paths, based on whether
-        # the input requires gradients and whether it lays on the CPU/GPU.
-        # For example, `torch_scatter` is usually faster than
-        # `torch.scatter_reduce` on GPU, while `torch.scatter_reduce` is faster
-        # on CPU.
-        # `torch.scatter_reduce` has a faster forward implementation for
-        # "min"/"max" reductions since it does not compute additional arg
-        # indices, but is therefore way slower in its backward implementation.
-        # More insights can be found in `test/utils/test_scatter.py`.
-
-        size = src.size()[:dim] + (dim_size, ) + src.size()[dim + 1:]
-
-        # For "any" reduction, we use regular `scatter_`:
-        if reduce == 'any':
-            index = broadcast(index, src, dim)
-            return src.new_zeros(size).scatter_(dim, index, src)
+warnings.filterwarnings('ignore', '.*is in beta and the API may change.*')
 
-        # For "sum" and "mean" reduction, we make use of `scatter_add_`:
-        if reduce == 'sum' or reduce == 'add':
-            index = broadcast(index, src, dim)
-            return src.new_zeros(size).scatter_add_(dim, index, src)
 
-        if reduce == 'mean':
-            count = src.new_zeros(dim_size)
-            count.scatter_add_(0, index, src.new_ones(src.size(dim)))
-            count = count.clamp(min=1)
+def scatter(
+    src: Tensor,
+    index: Tensor,
+    dim: int = 0,
+    dim_size: Optional[int] = None,
+    reduce: str = 'sum',
+) -> Tensor:
+    r"""Reduces all values from the :obj:`src` tensor at the indices specified
+    in the :obj:`index` tensor along a given dimension ``dim``. See the
+    `documentation <https://pytorch-scatter.readthedocs.io/en/latest/functions/scatter.html>`__  # noqa: E501
+    of the ``torch_scatter`` package for more information.
 
-            index = broadcast(index, src, dim)
-            out = src.new_zeros(size).scatter_add_(dim, index, src)
+    Args:
+        src (torch.Tensor): The source tensor.
+        index (torch.Tensor): The index tensor.
+        dim (int, optional): The dimension along which to index.
+            (default: ``0``)
+        dim_size (int, optional): The size of the output tensor at dimension
+            ``dim``. If set to :obj:`None`, will create a minimal-sized output
+            tensor according to ``index.max() + 1``. (default: :obj:`None`)
+        reduce (str, optional): The reduce operation (``"sum"``, ``"mean"``,
+            ``"mul"``, ``"min"``, ``"max"`` or ``"any"``). (default: ``"sum"``)
+    """
+    if isinstance(index, Tensor) and index.dim() != 1:
+        raise ValueError(f"The `index` argument must be one-dimensional "
+                         f"(got {index.dim()} dimensions)")
+
+    dim = src.dim() + dim if dim < 0 else dim
 
-            return out / broadcast(count, out, dim)
+    if isinstance(src, Tensor) and (dim < 0 or dim >= src.dim()):
+        raise ValueError(f"The `dim` argument must lay between 0 and "
+                         f"{src.dim() - 1} (got {dim})")
 
-        # For "min" and "max" reduction, we prefer `scatter_reduce_` on CPU or
-        # in case the input does not require gradients:
-        if reduce in ['min', 'max', 'amin', 'amax']:
-            if (not torch_geometric.typing.WITH_TORCH_SCATTER
-                    or is_compiling() or not src.is_cuda
-                    or not src.requires_grad):
+    if dim_size is None:
+        dim_size = int(index.max()) + 1 if index.numel() > 0 else 0
 
-                if src.is_cuda and src.requires_grad and not is_compiling():
-                    warnings.warn(f"The usage of `scatter(reduce='{reduce}')` "
-                                  f"can be accelerated via the 'torch-scatter'"
-                                  f" package, but it was not found")
+    # For now, we maintain various different code paths, based on whether
+    # the input requires gradients and whether it lays on the CPU/GPU.
+    # For example, `torch_scatter` is usually faster than
+    # `torch.scatter_reduce` on GPU, while `torch.scatter_reduce` is faster
+    # on CPU.
+    # `torch.scatter_reduce` has a faster forward implementation for
+    # "min"/"max" reductions since it does not compute additional arg
+    # indices, but is therefore way slower in its backward implementation.
+    # More insights can be found in `test/utils/test_scatter.py`.
+
+    size = src.size()[:dim] + (dim_size, ) + src.size()[dim + 1:]
+
+    # For "any" reduction, we use regular `scatter_`:
+    if reduce == 'any':
+        index = broadcast(index, src, dim)
+        return src.new_zeros(size).scatter_(dim, index, src)
+
+    # For "sum" and "mean" reduction, we make use of `scatter_add_`:
+    if reduce == 'sum' or reduce == 'add':
+        index = broadcast(index, src, dim)
+        return src.new_zeros(size).scatter_add_(dim, index, src)
+
+    if reduce == 'mean':
+        count = src.new_zeros(dim_size)
+        count.scatter_add_(0, index, src.new_ones(src.size(dim)))
+        count = count.clamp(min=1)
+
+        index = broadcast(index, src, dim)
+        out = src.new_zeros(size).scatter_add_(dim, index, src)
+
+        return out / broadcast(count, out, dim)
+
+    # For "min" and "max" reduction, we prefer `scatter_reduce_` on CPU or
+    # in case the input does not require gradients:
+    if reduce in ['min', 'max', 'amin', 'amax']:
+        if (not torch_geometric.typing.WITH_TORCH_SCATTER or is_compiling()
+                or is_in_onnx_export() or not src.is_cuda
+                or not src.requires_grad):
+
+            if (src.is_cuda and src.requires_grad and not is_compiling()
+                    and not is_in_onnx_export()):
+                warnings.warn(
+                    f"The usage of `scatter(reduce='{reduce}')` "
+                    f"can be accelerated via the 'torch-scatter'"
+                    f" package, but it was not found", stacklevel=2)
 
-                index = broadcast(index, src, dim)
+            index = broadcast(index, src, dim)
+            if not is_in_onnx_export():
                 return src.new_zeros(size).scatter_reduce_(
                     dim, index, src, reduce=f'a{reduce[-3:]}',
                     include_self=False)
 
-            return torch_scatter.scatter(src, index, dim, dim_size=dim_size,
-                                         reduce=reduce[-3:])
-
-        # For "mul" reduction, we prefer `scatter_reduce_` on CPU:
-        if reduce == 'mul':
-            if (not torch_geometric.typing.WITH_TORCH_SCATTER
-                    or is_compiling() or not src.is_cuda):
-
-                if src.is_cuda and not is_compiling():
-                    warnings.warn(f"The usage of `scatter(reduce='{reduce}')` "
-                                  f"can be accelerated via the 'torch-scatter'"
-                                  f" package, but it was not found")
-
-                index = broadcast(index, src, dim)
-                # We initialize with `one` here to match `scatter_mul` output:
-                return src.new_ones(size).scatter_reduce_(
-                    dim, index, src, reduce='prod', include_self=True)
-
-            return torch_scatter.scatter(src, index, dim, dim_size=dim_size,
-                                         reduce='mul')
-
-        raise ValueError(f"Encountered invalid `reduce` argument '{reduce}'")
-
-else:  # pragma: no cover
-
-    def scatter(
-        src: Tensor,
-        index: Tensor,
-        dim: int = 0,
-        dim_size: Optional[int] = None,
-        reduce: str = 'sum',
-    ) -> Tensor:
-        r"""Reduces all values from the :obj:`src` tensor at the indices
-        specified in the :obj:`index` tensor along a given dimension
-        :obj:`dim`. See the `documentation
-        <https://pytorch-scatter.readthedocs.io/en/latest/functions/
-        scatter.html>`_ of the :obj:`torch_scatter` package for more
-        information.
-
-        Args:
-            src (torch.Tensor): The source tensor.
-            index (torch.Tensor): The index tensor.
-            dim (int, optional): The dimension along which to index.
-                (default: :obj:`0`)
-            dim_size (int, optional): The size of the output tensor at
-                dimension :obj:`dim`. If set to :obj:`None`, will create a
-                minimal-sized output tensor according to
-                :obj:`index.max() + 1`. (default: :obj:`None`)
-            reduce (str, optional): The reduce operation (:obj:`"sum"`,
-                :obj:`"mean"`, :obj:`"mul"`, :obj:`"min"` or :obj:`"max"`,
-                :obj:`"any"`). (default: :obj:`"sum"`)
-        """
-        if reduce == 'any':
-            dim = src.dim() + dim if dim < 0 else dim
-
-            if dim_size is None:
-                dim_size = int(index.max()) + 1 if index.numel() > 0 else 0
-
-            size = src.size()[:dim] + (dim_size, ) + src.size()[dim + 1:]
+            fill = torch.full(  # type: ignore
+                size=(1, ),
+                fill_value=src.min() if 'max' in reduce else src.max(),
+                dtype=src.dtype,
+                device=src.device,
+            ).expand_as(src)
+            out = src.new_zeros(size).scatter_reduce_(dim, index, fill,
+                                                      reduce=f'a{reduce[-3:]}',
+                                                      include_self=True)
+            return out.scatter_reduce_(dim, index, src,
+                                       reduce=f'a{reduce[-3:]}',
+                                       include_self=True)
 
-            index = broadcast(index, src, dim)
-            return src.new_zeros(size).scatter_(dim, index, src)
+        return torch_scatter.scatter(src, index, dim, dim_size=dim_size,
+                                     reduce=reduce[-3:])
+
+    # For "mul" reduction, we prefer `scatter_reduce_` on CPU:
+    if reduce == 'mul':
+        if (not torch_geometric.typing.WITH_TORCH_SCATTER or is_compiling()
+                or not src.is_cuda):
 
-        if not torch_geometric.typing.WITH_TORCH_SCATTER:
-            raise ImportError("'scatter' requires the 'torch-scatter' package")
+            if src.is_cuda and not is_compiling():
+                warnings.warn(
+                    f"The usage of `scatter(reduce='{reduce}')` "
+                    f"can be accelerated via the 'torch-scatter'"
+                    f" package, but it was not found", stacklevel=2)
 
-        if reduce == 'amin' or reduce == 'amax':
-            reduce = reduce[-3:]
+            index = broadcast(index, src, dim)
+            # We initialize with `one` here to match `scatter_mul` output:
+            return src.new_ones(size).scatter_reduce_(dim, index, src,
+                                                      reduce='prod',
+                                                      include_self=True)
 
         return torch_scatter.scatter(src, index, dim, dim_size=dim_size,
-                                     reduce=reduce)
+                                     reduce='mul')
+
+    raise ValueError(f"Encountered invalid `reduce` argument '{reduce}'")
 
 
 def broadcast(src: Tensor, ref: Tensor, dim: int) -> Tensor:
@@ -187,7 +151,8 @@ def scatter_argmax(
     dim_size: Optional[int] = None,
 ) -> Tensor:
 
-    if torch_geometric.typing.WITH_TORCH_SCATTER and not is_compiling():
+    if (torch_geometric.typing.WITH_TORCH_SCATTER and not is_compiling()
+            and not is_in_onnx_export()):
         out = torch_scatter.scatter_max(src, index, dim=dim, dim_size=dim_size)
         return out[1]
 
@@ -199,15 +164,18 @@ def scatter_argmax(
     if dim_size is None:
         dim_size = int(index.max()) + 1 if index.numel() > 0 else 0
 
-    if torch_geometric.typing.WITH_PT112:
+    if not is_in_onnx_export():
         res = src.new_empty(dim_size)
         res.scatter_reduce_(0, index, src.detach(), reduce='amax',
                             include_self=False)
-    elif torch_geometric.typing.WITH_PT111:
-        res = torch.scatter_reduce(src.detach(), 0, index, reduce='amax',
-                                   output_size=dim_size)  # type: ignore
     else:
-        raise ValueError("'scatter_argmax' requires PyTorch >= 1.11")
+        # `include_self=False` is currently not supported by ONNX:
+        res = src.new_full(
+            size=(dim_size, ),
+            fill_value=src.min(),  # type: ignore
+        )
+        res.scatter_reduce_(0, index, src.detach(), reduce="amax",
+                            include_self=True)
 
     out = index.new_full((dim_size, ), fill_value=dim_size - 1)
     nonzero = (src == res[index]).nonzero().view(-1)
@@ -265,13 +233,7 @@ def group_argsort(
 
     # Compute `grouped_argsort`:
     src = src - 2 * index if descending else src + 2 * index
-    if torch_geometric.typing.WITH_PT113:
-        perm = src.argsort(descending=descending, stable=stable)
-    else:
-        perm = src.argsort(descending=descending)
-        if stable:
-            warnings.warn("Ignoring option `stable=True` in 'group_argsort' "
-                          "since it requires PyTorch >= 1.13.0")
+    perm = src.argsort(descending=descending, stable=stable)
     out = torch.empty_like(index)
     out[perm] = torch.arange(index.numel(), device=index.device)
 
@@ -295,7 +257,7 @@ def group_cat(
     r"""Concatenates the given sequence of tensors :obj:`tensors` in the given
     dimension :obj:`dim`.
     Different from :meth:`torch.cat`, values along the concatenating dimension
-    are grouped according to the indicies defined in the :obj:`index` tensors.
+    are grouped according to the indices defined in the :obj:`index` tensors.
     All tensors must have the same shape (except in the concatenating
     dimension).
 
@@ -326,5 +288,5 @@ def group_cat(
     """
     assert len(tensors) == len(indices)
     index, perm = torch.cat(indices).sort(stable=True)
-    out = torch.cat(tensors, dim=0)[perm]
+    out = torch.cat(tensors, dim=dim).index_select(dim, perm)
     return (out, index) if return_index else out

torch_geometric/utils/_sort_edge_index.py

@@ -107,8 +107,6 @@ def sort_edge_index(  # noqa: F811
     num_nodes = maybe_num_nodes(edge_index, num_nodes)
 
     if num_nodes * num_nodes > torch_geometric.typing.MAX_INT64:
-        if not torch_geometric.typing.WITH_PT113:
-            raise ValueError("'sort_edge_index' will result in an overflow")
         perm = lexsort(keys=[
             edge_index[int(sort_by_row)],
             edge_index[1 - int(sort_by_row)],

torch_geometric/utils/_spmm.py

@@ -63,18 +63,20 @@ def spmm(
 
         # Always convert COO to CSR for more efficient processing:
         if src.layout == torch.sparse_coo:
-            warnings.warn(f"Converting sparse tensor to CSR format for more "
-                          f"efficient processing. Consider converting your "
-                          f"sparse tensor to CSR format beforehand to avoid "
-                          f"repeated conversion (got '{src.layout}')")
+            warnings.warn(
+                f"Converting sparse tensor to CSR format for more "
+                f"efficient processing. Consider converting your "
+                f"sparse tensor to CSR format beforehand to avoid "
+                f"repeated conversion (got '{src.layout}')", stacklevel=2)
             src = src.to_sparse_csr()
 
         # Warn in case of CSC format without gradient computation:
         if src.layout == torch.sparse_csc and not other.requires_grad:
-            warnings.warn(f"Converting sparse tensor to CSR format for more "
-                          f"efficient processing. Consider converting your "
-                          f"sparse tensor to CSR format beforehand to avoid "
-                          f"repeated conversion (got '{src.layout}')")
+            warnings.warn(
+                f"Converting sparse tensor to CSR format for more "
+                f"efficient processing. Consider converting your "
+                f"sparse tensor to CSR format beforehand to avoid "
+                f"repeated conversion (got '{src.layout}')", stacklevel=2)
 
         # Use the default code path for `sum` reduction (works on CPU/GPU):
         if reduce == 'sum':
@@ -99,10 +101,11 @@ def spmm(
         # TODO The `torch.sparse.mm` code path with the `reduce` argument does
         # not yet support CSC :(
         if src.layout == torch.sparse_csc:
-            warnings.warn(f"Converting sparse tensor to CSR format for more "
-                          f"efficient processing. Consider converting your "
-                          f"sparse tensor to CSR format beforehand to avoid "
-                          f"repeated conversion (got '{src.layout}')")
+            warnings.warn(
+                f"Converting sparse tensor to CSR format for more "
+                f"efficient processing. Consider converting your "
+                f"sparse tensor to CSR format beforehand to avoid "
+                f"repeated conversion (got '{src.layout}')", stacklevel=2)
             src = src.to_sparse_csr()
 
         return torch.sparse.mm(src, other, reduce)
@@ -115,8 +118,7 @@ def spmm(
         if src.layout == torch.sparse_csr:
             ptr = src.crow_indices()
             deg = ptr[1:] - ptr[:-1]
-        elif (torch_geometric.typing.WITH_PT112
-              and src.layout == torch.sparse_csc):
+        elif src.layout == torch.sparse_csc:
             assert src.layout == torch.sparse_csc
             ones = torch.ones_like(src.values())
             index = src.row_indices()

torch_geometric/utils/_subgraph.py

@@ -346,10 +346,12 @@ def k_hop_subgraph(
 
     subsets = [node_idx]
 
+    preserved_edge_mask = torch.zeros_like(edge_mask)
     for _ in range(num_hops):
         node_mask.fill_(False)
         node_mask[subsets[-1]] = True
         torch.index_select(node_mask, 0, row, out=edge_mask)
+        preserved_edge_mask |= edge_mask
         subsets.append(col[edge_mask])
 
     subset, inv = torch.cat(subsets).unique(return_inverse=True)
@@ -360,6 +362,8 @@ def k_hop_subgraph(
 
     if not directed:
         edge_mask = node_mask[row] & node_mask[col]
+    else:
+        edge_mask = preserved_edge_mask
 
     edge_index = edge_index[:, edge_mask]
 

torch_geometric/utils/_tree_decomposition.py

@@ -64,7 +64,7 @@ def tree_decomposition(
             xs.append(1)
 
     # Generate `atom2cliques` mappings.
-    atom2cliques: List[List[int]] = [[] for i in range(mol.GetNumAtoms())]
+    atom2cliques: List[List[int]] = [[] for _ in range(mol.GetNumAtoms())]
     for c in range(len(cliques)):
         for atom in cliques[c]:
             atom2cliques[atom].append(c)

torch_geometric/utils/_trim_to_layer.py

@@ -234,10 +234,10 @@ def trim_sparse_tensor(src: SparseTensor, size: Tuple[int, int],
     rowptr = torch.narrow(rowptr, 0, 0, size[0] + 1).clone()
     rowptr[num_seed_nodes + 1:] = rowptr[num_seed_nodes]
 
-    col = torch.narrow(col, 0, 0, rowptr[-1])
+    col = torch.narrow(col, 0, 0, rowptr[-1])  # type: ignore
 
     if value is not None:
-        value = torch.narrow(value, 0, 0, rowptr[-1])
+        value = torch.narrow(value, 0, 0, rowptr[-1])  # type: ignore
 
     csr2csc = src.storage._csr2csc
     if csr2csc is not None:

torch_geometric/utils/augmentation.py

@@ -12,7 +12,7 @@ def shuffle_node(
     training: bool = True,
 ) -> Tuple[Tensor, Tensor]:
     r"""Randomly shuffle the feature matrix :obj:`x` along the
-    first dimmension.
+    first dimension.
 
     The method returns (1) the shuffled :obj:`x`, (2) the permutation
     indicating the orders of original nodes after shuffling.

torch_geometric/utils/convert.py

@@ -251,13 +251,13 @@ def from_networkx(
     if group_edge_attrs is not None and not isinstance(group_edge_attrs, list):
         group_edge_attrs = edge_attrs
 
-    for i, (_, feat_dict) in enumerate(G.nodes(data=True)):
+    for _, feat_dict in G.nodes(data=True):
         if set(feat_dict.keys()) != set(node_attrs):
             raise ValueError('Not all nodes contain the same attributes')
         for key, value in feat_dict.items():
             data_dict[str(key)].append(value)
 
-    for i, (_, _, feat_dict) in enumerate(G.edges(data=True)):
+    for _, _, feat_dict in G.edges(data=True):
         if set(feat_dict.keys()) != set(edge_attrs):
             raise ValueError('Not all edges contain the same attributes')
         for key, value in feat_dict.items():
@@ -452,15 +452,22 @@ def to_cugraph(
     g = cugraph.Graph(directed=directed)
     df = cudf.from_dlpack(to_dlpack(edge_index.t()))
 
+    df = cudf.DataFrame({
+        'source':
+        cudf.from_dlpack(to_dlpack(edge_index[0])),
+        'destination':
+        cudf.from_dlpack(to_dlpack(edge_index[1])),
+    })
+
     if edge_weight is not None:
         assert edge_weight.dim() == 1
-        df['2'] = cudf.from_dlpack(to_dlpack(edge_weight))
+        df['weight'] = cudf.from_dlpack(to_dlpack(edge_weight))
 
     g.from_cudf_edgelist(
         df,
-        source=0,
-        destination=1,
-        edge_attr='2' if edge_weight is not None else None,
+        source='source',
+        destination='destination',
+        edge_attr='weight' if edge_weight is not None else None,
         renumber=relabel_nodes,
     )
 
@@ -476,13 +483,13 @@ def from_cugraph(g: Any) -> Tuple[Tensor, Optional[Tensor]]:
     """
     df = g.view_edge_list()
 
-    src = from_dlpack(df[0].to_dlpack()).long()
-    dst = from_dlpack(df[1].to_dlpack()).long()
+    src = from_dlpack(df[g.source_columns].to_dlpack()).long()
+    dst = from_dlpack(df[g.destination_columns].to_dlpack()).long()
     edge_index = torch.stack([src, dst], dim=0)
 
     edge_weight = None
-    if '2' in df:
-        edge_weight = from_dlpack(df['2'].to_dlpack())
+    if g.weight_column is not None:
+        edge_weight = from_dlpack(df[g.weight_column].to_dlpack())
 
     return edge_index, edge_weight
 

torch_geometric/utils/cross_entropy.py

@@ -18,30 +18,51 @@ class SparseCrossEntropy(torch.autograd.Function):
     ) -> Tensor:
         assert inputs.dim() == 2
 
-        logsumexp = inputs.logsumexp(dim=-1)
-        ctx.save_for_backward(inputs, edge_label_index, edge_label_weight,
-                              logsumexp)
+        # Support for both positive and negative weights:
+        # Positive weights scale the logits *after* softmax.
+        # Negative weights scale the denominator *before* softmax:
+        pos_y = edge_label_index
+        neg_y = pos_weight = neg_weight = None
 
-        out = inputs[edge_label_index[0], edge_label_index[1]]
-        out.neg_().add_(logsumexp[edge_label_index[0]])
         if edge_label_weight is not None:
-            out *= edge_label_weight
+            pos_mask = edge_label_weight >= 0
+            pos_y = edge_label_index[:, pos_mask]
+            pos_weight = edge_label_weight[pos_mask]
+
+            if pos_y.size(1) < edge_label_index.size(1):
+                neg_mask = ~pos_mask
+                neg_y = edge_label_index[:, neg_mask]
+                neg_weight = edge_label_weight[neg_mask]
+
+            if neg_y is not None and neg_weight is not None:
+                inputs = inputs.clone()
+                inputs[
+                    neg_y[0],
+                    neg_y[1],
+                ] += neg_weight.abs().log().clamp(min=1e-12)
+
+        logsumexp = inputs.logsumexp(dim=-1)
+        ctx.save_for_backward(inputs, pos_y, pos_weight, logsumexp)
+
+        out = inputs[pos_y[0], pos_y[1]]
+        out.neg_().add_(logsumexp[pos_y[0]])
+        if pos_weight is not None:
+            out *= pos_weight
 
         return out.sum() / inputs.size(0)
 
     @staticmethod
     @torch.autograd.function.once_differentiable
     def backward(ctx: Any, grad_out: Tensor) -> Tuple[Tensor, None, None]:
-        inputs, edge_label_index, edge_label_weight, logsumexp = (
-            ctx.saved_tensors)
+        inputs, pos_y, pos_weight, logsumexp = ctx.saved_tensors
 
         grad_out = grad_out / inputs.size(0)
-        grad_out = grad_out.expand(edge_label_index.size(1))
+        grad_out = grad_out.expand(pos_y.size(1))
 
-        if edge_label_weight is not None:
-            grad_out = grad_out * edge_label_weight
+        if pos_weight is not None:
+            grad_out = grad_out * pos_weight
 
-        grad_logsumexp = scatter(grad_out, edge_label_index[0], dim=0,
+        grad_logsumexp = scatter(grad_out, pos_y[0], dim=0,
                                  dim_size=inputs.size(0), reduce='sum')
 
         # Gradient computation of `logsumexp`: `grad * (self - result).exp()`
@@ -49,7 +70,7 @@ class SparseCrossEntropy(torch.autograd.Function):
         grad_input.exp_()
         grad_input.mul_(grad_logsumexp.view(-1, 1))
 
-        grad_input[edge_label_index[0], edge_label_index[1]] -= grad_out
+        grad_input[pos_y[0], pos_y[1]] -= grad_out
 
         return grad_input, None, None