Repository: dmlc/dgl

Issue:
multi_update_all gives infinity on isolated nodes when using min reducer with different output keys
## 🐛 Bug

The title is quite long so I will directly demonstrate the bug below.

## To Reproduce

```python
import dgl
import torch
g = dgl.heterograph({
    ('A', 'AB', 'B'): ([0,1,2,3],[0,0,1,1]),
    ('C', 'CB', 'B'): ([0,1,2,3],[2,2,3,3])
})
g.nodes['A'].data['x'] = torch.randn(4, 16)
g.nodes['C'].data['x'] = torch.randn(4, 16)
g.multi_update_all({
    'AB': (dgl.function.copy_u('x', 'm'), dgl.function.min('m', 'a1')),
    'CB': (dgl.function.copy_u('x', 'm'), dgl.function.min('m', 'a2'))
    }, cross_reducer='min')
# One of the following gives infinity while the other does not.
print(g.nodes['B'].data['a1'])
print(g.nodes['B'].data['a2'])
```

## Expected behavior

Neither of the outputs should give infinity.

## Environment

 - DGL Version (e.g., 1.0): master
 - Backend Library & Version (e.g., PyTorch 0.4.1, MXNet/Gluon 1.3):
 - OS (e.g., Linux):
 - How you installed DGL (`conda`, `pip`, source):
 - Build command you used (if compiling from source):
 - Python version:
 - CUDA/cuDNN version (if applicable):
 - GPU models and configuration (e.g. V100):
 - Any other relevant information:

## Additional context

The reason is that in https://github.com/dmlc/dgl/blob/3a0bcc0c15cd0ce445318382887cd977979514df/python/dgl/heterograph.py#L5261-L5271
we are only replacing infinities for the first output key, assuming that the reducers all produce the same output key.  In practice we should use a for loop.

Credit to @zhhhhahahaha for finding this bug.



Discussion:

Resolving patch:
diff --git a/python/dgl/heterograph.py b/python/dgl/heterograph.py
--- a/python/dgl/heterograph.py
+++ b/python/dgl/heterograph.py
@@ -5263,12 +5263,12 @@ def multi_update_all(self, etype_dict, cross_reducer, apply_node_func=None):
             out = reduce_dict_data(frames, cross_reducer, merge_order[dtid])
             # Replace infinity with zero for isolated nodes when reducer is min/max
             if core.is_builtin(rfunc) and rfunc.name in ["min", "max"]:
-                key = list(out.keys())[0]
-                out[key] = (
-                    F.replace_inf_with_zero(out[key])
-                    if out[key] is not None
-                    else None
-                )
+                for key in out.keys():
+                    out[key] = (
+                        F.replace_inf_with_zero(out[key])
+                        if out[key] is not None
+                        else None
+                    )
             self._node_frames[dtid].update(out)
             # apply
             if apply_node_func is not None:


Test patch:
diff --git a/tests/python/common/test_heterograph-update-all.py b/tests/python/common/test_heterograph-update-all.py
--- a/tests/python/common/test_heterograph-update-all.py
+++ b/tests/python/common/test_heterograph-update-all.py
@@ -334,6 +334,38 @@ def _print_error(a, b):
                 _test(lhs, rhs, binary_op, reducer)
 
 
+# Issue #5873
+def test_multi_update_all_minmax_reduce_with_isolated_nodes():
+    g = dgl.heterograph(
+        {
+            ("A", "AB", "B"): ([0, 1, 2, 3], [0, 0, 1, 1]),
+            ("C", "CB", "B"): ([0, 1, 2, 3], [2, 2, 3, 3]),
+        },
+        device=F.ctx(),
+    )
+    g.nodes["A"].data["x"] = F.randn((4, 16))
+    g.nodes["C"].data["x"] = F.randn((4, 16))
+    g.multi_update_all(
+        {
+            "AB": (dgl.function.copy_u("x", "m"), dgl.function.min("m", "a1")),
+            "CB": (dgl.function.copy_u("x", "m"), dgl.function.min("m", "a2")),
+        },
+        cross_reducer="min",
+    )
+    assert not np.isinf(F.asnumpy(g.nodes["B"].data["a1"])).any()
+    assert not np.isinf(F.asnumpy(g.nodes["B"].data["a2"])).any()
+
+    g.multi_update_all(
+        {
+            "AB": (dgl.function.copy_u("x", "m"), dgl.function.max("m", "a1")),
+            "CB": (dgl.function.copy_u("x", "m"), dgl.function.max("m", "a2")),
+        },
+        cross_reducer="max",
+    )
+    assert not np.isinf(F.asnumpy(g.nodes["B"].data["a1"])).any()
+    assert not np.isinf(F.asnumpy(g.nodes["B"].data["a2"])).any()
+
+
 if __name__ == "__main__":
     test_unary_copy_u()
     test_unary_copy_e()