Repository: pyg-team/pytorch_geometric
Pull request #7334 (2023-05-09 21:58:54)

[Aggregation] Clamp `StdAggregation` if there are edges for aggregation
If there are no edges for aggregation and we specified `dim_size`, IMO std should be all zeros.
Previous implementation will clamp value to 0.00001 and in this scenario the aggregated std value will be around 0.0032.

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

> :exclamation: Current head 9f13b17 differs from pull request most recent head 80d11ed. Consider uploading reports for the commit 80d11ed to get more accurate results

```diff
@@           Coverage Diff           @@
##           master    #7334   +/-   ##
=======================================
  Coverage   91.52%   91.53%           
=======================================
  Files         438      438           
  Lines       24355    24357    +2     
=======================================
+ Hits        22292    22294    +2     
  Misses       2063     2063           
```


| [Impacted Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7334?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/nn/aggr/basic.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7334?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL2FnZ3IvYmFzaWMucHk=) | `95.87% <100.00%> (+0.04%)` | :arrow_up: |
| [torch\_geometric/nn/aggr/fused.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7334?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL2FnZ3IvZnVzZWQucHk=) | `97.40% <100.00%> (+0.01%)` | :arrow_up: |

: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/torch_geometric/nn/aggr/basic.py b/torch_geometric/nn/aggr/basic.py
index e3a93211d..324d4a7b4 100644
--- a/torch_geometric/nn/aggr/basic.py
+++ b/torch_geometric/nn/aggr/basic.py
@@ -1,3 +1,4 @@
+import math
 from typing import Optional
 
 import torch
@@ -131,7 +132,10 @@ class StdAggregation(Aggregation):
                 ptr: Optional[Tensor] = None, dim_size: Optional[int] = None,
                 dim: int = -2) -> Tensor:
         var = self.var_aggr(x, index, ptr, dim_size, dim)
-        return var.clamp(min=1e-5).sqrt()
+        # Allow "undefined" gradient at `sqrt(0.0)`:
+        out = var.clamp(min=1e-5).sqrt()
+        out = out.masked_fill(out <= math.sqrt(1e-5), 0.0)
+        return out
 
 
 class SoftmaxAggregation(Aggregation):
diff --git a/torch_geometric/nn/aggr/fused.py b/torch_geometric/nn/aggr/fused.py
index 965ac6517..97a00aad9 100644
--- a/torch_geometric/nn/aggr/fused.py
+++ b/torch_geometric/nn/aggr/fused.py
@@ -1,3 +1,4 @@
+import math
 from typing import Dict, List, Optional, Tuple, Union
 
 from torch import Tensor
@@ -316,7 +317,11 @@ class FusedAggregation(Aggregation):
                 assert mean is not None
                 var = (pow_sum / count) - (mean * mean)
 
-            outs[i] = var.clamp(min=1e-5).sqrt()
+            # Allow "undefined" gradient at `sqrt(0.0)`:
+            out = var.clamp(min=1e-5).sqrt()
+            out = out.masked_fill(out <= math.sqrt(1e-5), 0.0)
+
+            outs[i] = out
 
         #######################################################################
 


Test patch:
diff --git a/test/nn/aggr/test_basic.py b/test/nn/aggr/test_basic.py
index 45d04dc3f..07841ec20 100644
--- a/test/nn/aggr/test_basic.py
+++ b/test/nn/aggr/test_basic.py
@@ -74,6 +74,17 @@ def test_var_aggregation():
     assert torch.allclose(out, expected, atol=1e-6)
 
 
+def test_empty_std_aggregation():
+    aggr = StdAggregation()
+
+    x = torch.empty(0, 6).reshape(0, 6)
+    index = torch.empty(0, dtype=torch.long)
+
+    out = aggr(x, index, dim_size=5)
+    assert out.size() == (5, 6)
+    assert float(out.abs().sum()) == 0.0
+
+
 @pytest.mark.parametrize('Aggregation', [
     SoftmaxAggregation,
     PowerMeanAggregation,
diff --git a/test/nn/aggr/test_fused.py b/test/nn/aggr/test_fused.py
index b6fea282a..b5a376283 100644
--- a/test/nn/aggr/test_fused.py
+++ b/test/nn/aggr/test_fused.py
@@ -30,16 +30,28 @@ def test_fused_aggregation(aggrs):
     out = torch.cat(aggr(x, index), dim=-1)
 
     expected = torch.cat([aggr(y, index) for aggr in aggrs], dim=-1)
-    assert torch.allclose(out, expected, atol=1e-6)
+    assert torch.allclose(out, expected, atol=1e-5)
 
     jit = torch.jit.script(aggr)
-    assert torch.allclose(torch.cat(jit(x, index), dim=-1), out, atol=1e-6)
+    assert torch.allclose(torch.cat(jit(x, index), dim=-1), out, atol=1e-5)
 
     out.mean().backward()
     assert x.grad is not None
     expected.mean().backward()
     assert y.grad is not None
-    assert torch.allclose(x.grad, y.grad)
+    assert torch.allclose(x.grad, y.grad, atol=1e-5)
+
+
+def test_empty_fused_std_aggregation():
+    aggrs = [aggregation_resolver(aggr) for aggr in ['mean', 'var', 'std']]
+    aggr = FusedAggregation(aggrs)
+
+    x = torch.empty(0, 6).reshape(0, 6)
+    index = torch.empty(0, dtype=torch.long)
+
+    out = torch.cat(aggr(x, index, dim_size=5), dim=-1)
+    assert out.size() == (5, 18)
+    assert float(out.abs().sum()) == 0.0
 
 
 if __name__ == '__main__':