Repository: OpenMined/PySyft

Issue:
Implement Default scatter Functionality for Base Tensor Type
**User Story A:** As a Data Scientist using Syft's Base Tensor type, I want to leverage a default method for computing operations on a Tensor of arbitrary type. For this ticket to be complete, scatter_() should operate inline. For a reference on the operation this performs check out [PyTorch](http://pytorch.org/docs/master/tensors.html)'s documentation.

**Acceptance Criteria:**
- If the Base Tensor type's attribute "encrypted" is set to True, it should return a NotImplemented error.
- a unit test demonstrating the correct operation on the Base Tensor type implemented over int and float Tensors.
- inline documentation in the python code. For inspiration on inline documentation, please check out [PyTorch](http://pytorch.org/docs/master/tensors.html)'s documentation for this operator.

Discussion:

Patch:
diff --git a/syft/tensor.py b/syft/tensor.py
index 0aed1c3ef8..7d31e2e430 100644
--- a/syft/tensor.py
+++ b/syft/tensor.py
@@ -1034,6 +1034,67 @@ class TensorBase(object):
         hist, edges = np.histogram(np.array(self.data), bins=bins, range=(min, max))
         return TensorBase(hist)
 
+    def scatter_(self, dim, index, src):
+        """
+        Writes all values from the Tensor src into self at the indices specified in the index Tensor.
+        The indices are specified with respect to the given dimension, dim, in the manner described in gather().
+
+        :param dim: The axis along which to index
+        :param index: The indices of elements to scatter
+        :param src: The source element(s) to scatter
+        :return: self
+        """
+        index = _ensure_tensorbase(index)
+        if self.encrypted or index.encrypted:
+            return NotImplemented
+        if index.data.dtype != np.dtype('int_'):
+            raise TypeError("The values of index must be integers")
+        if self.data.ndim != index.data.ndim:
+            raise ValueError("Index should have the same number of dimensions as output")
+        if dim >= self.data.ndim or dim < -self.data.ndim:
+            raise IndexError("dim is out of range")
+        if dim < 0:
+            # Not sure why scatter should accept dim < 0, but that is the behavior in PyTorch's scatter
+            dim = self.data.ndim + dim
+        idx_xsection_shape = index.data.shape[:dim] + index.data.shape[dim + 1:]
+        self_xsection_shape = self.data.shape[:dim] + self.data.shape[dim + 1:]
+        if idx_xsection_shape != self_xsection_shape:
+            raise ValueError("Except for dimension " + str(dim) +
+                             ", all dimensions of index and output should be the same size")
+        if (index.data >= self.data.shape[dim]).any() or (index.data < 0).any():
+            raise IndexError("The values of index must be between 0 and (self.data.shape[dim] -1)")
+
+        def make_slice(arr, dim, i):
+            slc = [slice(None)] * arr.ndim
+            slc[dim] = i
+            return slc
+
+        # We use index and dim parameters to create idx
+        # idx is in a form that can be used as a NumPy advanced index for scattering of src param. in self.data
+        idx = [[*np.indices(idx_xsection_shape).reshape(index.data.ndim - 1, -1),
+                index.data[make_slice(index.data, dim, i)].reshape(1, -1)[0]] for i in range(index.data.shape[dim])]
+        idx = list(np.concatenate(idx, axis=1))
+        idx.insert(dim, idx.pop())
+
+        if not np.isscalar(src):
+            src = _ensure_tensorbase(src)
+            if index.data.shape[dim] > src.data.shape[dim]:
+                raise IndexError("Dimension " + str(dim) + "of index can not be bigger than that of src ")
+            src_shape = src.data.shape[:dim] + src.data.shape[dim + 1:]
+            if idx_xsection_shape != src_shape:
+                raise ValueError("Except for dimension " +
+                                 str(dim) + ", all dimensions of index and src should be the same size")
+            # src_idx is a NumPy advanced index for indexing of elements in the src
+            src_idx = list(idx)
+            src_idx.pop(dim)
+            src_idx.insert(dim, np.repeat(np.arange(index.data.shape[dim]), np.prod(idx_xsection_shape)))
+            self.data[idx] = src.data[src_idx]
+
+        else:
+            self.data[idx] = src
+
+        return self
+
     def serialize(self):
         return pickle.dumps(self)
 


Test patch:
diff --git a/tests/test_tensor.py b/tests/test_tensor.py
index d666b46e88..f586ff59ac 100644
--- a/tests/test_tensor.py
+++ b/tests/test_tensor.py
@@ -726,5 +726,100 @@ class notEqualTests(unittest.TestCase):
         self.assertTrue(syft.equal(t1, TensorBase([1, 1, 1, 0])))
 
 
+class scatterTests(unittest.TestCase):
+    def testScatter_Numerical0(self):
+        t = TensorBase(np.zeros((3, 5)))
+        idx = TensorBase(np.array([[0, 0, 0, 0, 0]]))
+        src = 1.0
+        dim = 0
+        t.scatter_(dim=dim, index=idx, src=src)
+        self.assertTrue(np.array_equal(t.data, np.array([[1, 1, 1, 1, 1], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0]])))
+
+    def testScatter_Numerical1(self):
+        t = TensorBase(np.zeros((3, 5)))
+        idx = TensorBase(np.array([[0], [0], [0]]))
+        src = 1.0
+        dim = 1
+        t.scatter_(dim=dim, index=idx, src=src)
+        self.assertTrue(np.array_equal(t.data, np.array([[1, 0, 0, 0, 0], [1, 0, 0, 0, 0], [1, 0, 0, 0, 0]])))
+
+    def testScatter_Numerical2(self):
+        t = TensorBase(np.zeros((3, 5)))
+        idx = TensorBase(np.array([[0], [0], [0]]))
+        src = 1.0
+        dim = -1
+        t.scatter_(dim=dim, index=idx, src=src)
+        self.assertTrue(np.array_equal(t.data, np.array([[1, 0, 0, 0, 0], [1, 0, 0, 0, 0], [1, 0, 0, 0, 0]])))
+
+    def testScatter_Numerical3(self):
+        t = TensorBase(np.zeros((3, 5)))
+        idx = TensorBase(np.array([[0, 0, 0, 0, 0]]))
+        src = TensorBase(np.array([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]))
+        dim = 0
+        t.scatter_(dim=dim, index=idx, src=src)
+        self.assertTrue(np.array_equal(t.data, np.array([[1, 2, 3, 4, 5], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0]])))
+
+    def testScatter_Numerical4(self):
+        t = TensorBase(np.zeros((3, 5)))
+        idx = TensorBase(np.array([[0, 0, 0, 0, 0]]))
+        src = TensorBase(np.array([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]))
+        dim = -2
+        t.scatter_(dim=dim, index=idx, src=src)
+        self.assertTrue(np.array_equal(t.data, np.array([[1, 2, 3, 4, 5], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0]])))
+
+    def testScatter_Numerical5(self):
+        t = TensorBase(np.zeros((3, 5)))
+        idx = TensorBase(np.array([[0, 0, 0, 0, 0], [0, 0, 0, 0, 0]]))
+        src = TensorBase(np.array([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]))
+        dim = 0
+        t.scatter_(dim=dim, index=idx, src=src)
+        self.assertTrue(np.array_equal(t.data, np.array([[6, 7, 8, 9, 10], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0]])))
+
+    def testScatter_Numerical6(self):
+        t = TensorBase(np.zeros((3, 4, 5)))
+        idx = [[[3, 0, 1, 1, 2], [0, 3, 3, 3, 3]], [[2, 0, 0, 0, 0], [2, 1, 0, 2, 0]],
+               [[0, 0, 1, 0, 2], [1, 3, 2, 2, 2]]]
+        src = [[[7, 84, 99, 71, 44], [79, 57, 2, 37, 62]], [[31, 44, 43, 54, 56], [72, 52, 21, 89, 95]],
+               [[5, 3, 99, 4, 52], [32, 88, 58, 62, 9]]]
+        dim = 1
+        t.scatter_(dim=dim, index=idx, src=src)
+        expected = [[[79, 84, 0, 0, 0], [0, 0, 99, 71, 0], [0, 0, 0, 0, 44], [7, 57, 2, 37, 62]],
+                    [[0, 44, 21, 54, 95], [0, 52, 0, 0, 0], [72, 0, 0, 89, 0], [0, 0, 0, 0, 0]],
+                    [[5, 3, 0, 4, 0], [32, 0, 99, 0, 0], [0, 0, 58, 62, 9], [0, 88, 0, 0, 0]]]
+        self.assertTrue(np.array_equal(t.data, np.array(expected)))
+
+    def testScatter_IndexType(self):
+        t = TensorBase(np.zeros((3, 5)))
+        idx = TensorBase(np.array([[0.0, 0.0, 0.0, 0.0, 0.0]]))
+        src = TensorBase(np.array([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]))
+        dim = 0
+        with self.assertRaises(Exception):
+            t.scatter_(dim=dim, index=idx, src=src)
+
+    def testScatter_IndexOutOfRange(self):
+        t = TensorBase(np.zeros((3, 5)))
+        idx = TensorBase(np.array([[5, 0, 0, 0, 0]]))
+        src = TensorBase(np.array([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]))
+        dim = 0
+        with self.assertRaises(Exception):
+            t.scatter_(dim=dim, index=idx, src=src)
+
+    def testScatter_DimOutOfRange(self):
+        t = TensorBase(np.zeros((3, 5)))
+        idx = TensorBase(np.array([[0, 0, 0, 0, 0]]))
+        src = TensorBase(np.array([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]))
+        dim = 4
+        with self.assertRaises(Exception):
+            t.scatter_(dim=dim, index=idx, src=src)
+
+    def testScatter_index_src_dimension_mismatch(self):
+        t = TensorBase(np.zeros((3, 5)))
+        idx = TensorBase(np.array([[0, 0, 0, 0, 0], [0, 0, 0, 0, 0]]))
+        src = TensorBase(np.array([[1, 2, 3, 4, 5]]))
+        dim = 1
+        with self.assertRaises(Exception):
+            t.scatter_(dim=dim, index=idx, src=src)
+
+
 if __name__ == "__main__":
     unittest.main()