Repository: pyg-team/pytorch_geometric
Pull request #9447 (2024-06-24 09:50:33)

Add `radd` and `rsub` in `Index`


Patch:
diff --git a/torch_geometric/index.py b/torch_geometric/index.py
index d0c5d7a81..af64bf396 100644
--- a/torch_geometric/index.py
+++ b/torch_geometric/index.py
@@ -685,14 +685,14 @@ def _index(
 
 @implements(aten.add.Tensor)
 def _add(
-    input: Index,
+    input: Union[int, Tensor, Index],
     other: Union[int, Tensor, Index],
     *,
     alpha: int = 1,
 ) -> Union[Index, Tensor]:
 
     data = aten.add.Tensor(
-        input._data,
+        input._data if isinstance(input, Index) else input,
         other._data if isinstance(other, Index) else other,
         alpha=alpha,
     )
@@ -704,15 +704,25 @@ def _add(
 
     out = Index(data)
 
+    if isinstance(input, Tensor) and input.numel() <= 1:
+        input = int(input)
+
     if isinstance(other, Tensor) and other.numel() <= 1:
         other = int(other)
 
     if isinstance(other, int):
+        assert isinstance(input, Index)
         if input.dim_size is not None:
             out._dim_size = input.dim_size + alpha * other
         out._is_sorted = input.is_sorted
 
-    elif isinstance(other, Index):
+    elif isinstance(input, int):
+        assert isinstance(other, Index)
+        if other.dim_size is not None:
+            out._dim_size = input + alpha * other.dim_size
+        out._is_sorted = other.is_sorted
+
+    elif isinstance(input, Index) and isinstance(other, Index):
         if input.dim_size is not None and other.dim_size is not None:
             out._dim_size = input.dim_size + alpha * other.dim_size
 
@@ -754,14 +764,14 @@ def add_(
 
 @implements(aten.sub.Tensor)
 def _sub(
-    input: Index,
+    input: Union[int, Tensor, Index],
     other: Union[int, Tensor, Index],
     *,
     alpha: int = 1,
 ) -> Union[Index, Tensor]:
 
     data = aten.sub.Tensor(
-        input._data,
+        input._data if isinstance(input, Index) else input,
         other._data if isinstance(other, Index) else other,
         alpha=alpha,
     )
@@ -773,6 +783,9 @@ def _sub(
 
     out = Index(data)
 
+    if not isinstance(input, Index):
+        return out
+
     if isinstance(other, Tensor) and other.numel() <= 1:
         other = int(other)
 


Test patch:
diff --git a/test/test_index.py b/test/test_index.py
index 5a20d795e..93a329b6c 100644
--- a/test/test_index.py
+++ b/test/test_index.py
@@ -446,6 +446,12 @@ def test_add(dtype, device):
     assert out.dim_size == 5
     assert out.is_sorted
 
+    out = tensor([2], dtype=dtype, device=device) + index
+    assert isinstance(out, Index)
+    assert out.equal(tensor([2, 3, 3, 4], device=device))
+    assert out.dim_size == 5
+    assert out.is_sorted
+
     out = index.add(index)
     assert isinstance(out, Index)
     assert out.equal(tensor([0, 2, 2, 4], device=device))
@@ -480,6 +486,12 @@ def test_sub(dtype, device):
     assert out.dim_size == 5
     assert out.is_sorted
 
+    out = tensor([6], dtype=dtype, device=device) - index
+    assert isinstance(out, Index)
+    assert out.equal(tensor([2, 1, 1, 0], device=device))
+    assert out.dim_size is None
+    assert not out.is_sorted
+
     out = index.sub(index)
     assert isinstance(out, Index)
     assert out.equal(tensor([0, 0, 0, 0], device=device))