Repository: pyg-team/pytorch_geometric
Pull request #7447 (2023-05-27 11:13:34)

Sparse `cross_entropy` implementation

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7447?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#7447](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7447?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (428cc7f) into [master](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/9e06952e5925cee9daa0ffcfb061ecd923621f20?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (9e06952) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

```diff
@@           Coverage Diff           @@
##           master    #7447   +/-   ##
=======================================
  Coverage   91.42%   91.42%           
=======================================
  Files         441      442    +1     
  Lines       24651    24658    +7     
=======================================
+ Hits        22536    22543    +7     
  Misses       2115     2115           
```


| [Impacted Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7447?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/utils/cross\_entropy.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7447?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3V0aWxzL2Nyb3NzX2VudHJvcHkucHk=) | `100.00% <100.00%> (ø)` | |

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 03674baac..f3a395972 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added a sparse `cross_entropy` implementation ([#7447](https://github.com/pyg-team/pytorch_geometric/pull/7447))
 - Added the `MovieLens-100K` heterogeneous dataset ([#7398](https://github.com/pyg-team/pytorch_geometric/pull/7398))
 - Added the `PMLP` model ([#7370](https://github.com/pyg-team/pytorch_geometric/pull/7370))
 - Added padding capabilities to `HeteroData.to_homogeneous()` in case feature dimensionalities do not match ([#7374](https://github.com/pyg-team/pytorch_geometric/pull/7374))
diff --git a/torch_geometric/utils/cross_entropy.py b/torch_geometric/utils/cross_entropy.py
new file mode 100644
index 000000000..a05793892
--- /dev/null
+++ b/torch_geometric/utils/cross_entropy.py
@@ -0,0 +1,29 @@
+from torch import Tensor
+
+
+def sparse_cross_entropy(inputs: Tensor, edge_label_index: Tensor) -> Tensor:
+    r"""A sparse-label variant of :func:`torch.nn.functional.cross_entropy`.
+    In particular, the binary target matrix is solely given by sparse indices
+    :obj:`edge_label_index`.
+
+    Args:
+        inputs (torch.Tensor): The predicted unnormalized logits of shape
+            :obj:`[batch_size, num_classes]`.
+        edge_index (torch.Tensor): The sparse ground-truth indices of
+            shape :obj:`[2, num_labels]`.
+
+    :rtype: :class:`torch.Tensor`
+
+    Example:
+
+        >>> inputs = torch.randn(2, 3)
+        >>> edge_label_index = torch.tensor([[0, 0, 1],
+        ...                                  [0, 1, 2]])
+        >>> sparse_cross_entropy(inputs, edge_label_index)
+        tensor(1.2919)
+    """
+    assert inputs.dim() == 2
+    logsumexp = inputs.logsumexp(dim=-1)
+    values = inputs[edge_label_index[0], edge_label_index[1]]
+    out = -values + logsumexp[edge_label_index[0]]
+    return out.sum() / inputs.size(0)


Test patch:
diff --git a/test/utils/test_cross_entropy.py b/test/utils/test_cross_entropy.py
new file mode 100644
index 000000000..c7fb1c7bd
--- /dev/null
+++ b/test/utils/test_cross_entropy.py
@@ -0,0 +1,24 @@
+import torch
+import torch.nn.functional as F
+
+from torch_geometric.utils.cross_entropy import sparse_cross_entropy
+
+
+def test_sparse_cross_entropy_multiclass():
+    x = torch.randn(5, 5)
+    y = torch.eye(5)
+    edge_label_index = y.nonzero().t()
+
+    expected = F.cross_entropy(x, y)
+    out = sparse_cross_entropy(x, edge_label_index)
+    assert torch.allclose(expected, out)
+
+
+def test_sparse_cross_entropy_multilabel():
+    x = torch.randn(5, 8)
+    y = torch.randint_like(x, 0, 2)
+    edge_label_index = y.nonzero().t()
+
+    expected = F.cross_entropy(x, y)
+    out = sparse_cross_entropy(x, edge_label_index)
+    assert torch.allclose(expected, out)