autograd: return self for += and -=. (OpenMined#3872)

Co-authored-by: Justin Harris <justindharris@gmail.com>

syft/frameworks/torch/tensors/interpreters/autograd.py

@@ -90,6 +90,7 @@ def __iadd__(self, other):
         result = self.add(other)
         self.child = result.child
         self.grad_fn = result.grad_fn
+        return self
 
     def __sub__(self, other):
         if isinstance(self, AutogradTensor) and not isinstance(other, AutogradTensor):
@@ -100,6 +101,7 @@ def __isub__(self, other):
         result = self.sub(other)
         self.child = result.child
         self.grad_fn = result.grad_fn
+        return self
 
     def __mul__(self, other):
         if isinstance(self, AutogradTensor) and not isinstance(other, AutogradTensor):

test/generic/test_autograd.py

@@ -22,6 +22,34 @@ def test_wrap():
     assert isinstance(x.child.child, torch.Tensor)
 
 
+def test_iadd():
+    """
+    Test += AutogradTensor
+    """
+
+    a = AutogradTensor(torch.Tensor([1, 2, 3]))
+    b = AutogradTensor(torch.Tensor([4, 5, 6]))
+    expected = AutogradTensor(torch.Tensor([5, 7, 9]))
+
+    a += b
+    assert a is not None
+    assert torch.equal(a, expected)
+
+
+def test_isub():
+    """
+    Test -= AutogradTensor
+    """
+
+    a = AutogradTensor(torch.Tensor([1, 2, 3]))
+    b = AutogradTensor(torch.Tensor([4, 5, 6]))
+    expected = AutogradTensor(torch.Tensor([-3, -3, -3]))
+
+    a -= b
+    assert a is not None
+    assert torch.equal(a, expected)
+
+
 @pytest.mark.parametrize("cmd", ["__add__", "__sub__", "__mul__", "__matmul__"])
 @pytest.mark.parametrize("backward_one", [True, False])
 def test_backward_for_binary_cmd_with_autograd(cmd, backward_one):