Repository: OpenMined/PySyft

Issue:
Implement Default gather Functionality for Base Tensor Type
**User Story A:** As a Data Scientist using Syft's Base Tensor type, we want to implement a default method for computing operations on a Tensor of arbitrary type. For this ticket to be complete, gather() should return a new tensor. 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 765d418508..a016eb8cfa 100644
--- a/syft/tensor.py
+++ b/syft/tensor.py
@@ -1051,9 +1051,8 @@ class TensorBase(object):
 
     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().
-
+        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
@@ -1110,6 +1109,32 @@ class TensorBase(object):
 
         return self
 
+    def gather(self, dim, index):
+        """
+        Gathers values along an axis specified by ``dim``.
+        For a 3-D tensor the output is specified by:
+            out[i][j][k] = input[index[i][j][k]][j][k]  # if dim == 0
+            out[i][j][k] = input[i][index[i][j][k]][k]  # if dim == 1
+            out[i][j][k] = input[i][j][index[i][j][k]]  # if dim == 2
+        :param dim: The axis along which to index
+        :param index: A tensor of indices of elements to gather
+        :return: tensor of gathered values
+        """
+        index = _ensure_tensorbase(index)
+        if self.encrypted or index.encrypted:
+            return NotImplemented
+        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 self should be the same size")
+        if index.data.dtype != np.dtype('int_'):
+            raise TypeError("The values of index must be integers")
+        data_swaped = np.swapaxes(self.data, 0, dim)
+        index_swaped = np.swapaxes(index, 0, dim)
+        gathered = np.choose(index_swaped, data_swaped)
+        return TensorBase(np.swapaxes(gathered, 0, dim))
+
     def serialize(self):
         return pickle.dumps(self)
 


Test patch:
diff --git a/tests/test_tensor.py b/tests/test_tensor.py
index ce477234fe..fbeb064be6 100644
--- a/tests/test_tensor.py
+++ b/tests/test_tensor.py
@@ -750,6 +750,23 @@ class notEqualTests(unittest.TestCase):
         self.assertTrue(syft.equal(t1, TensorBase([1, 1, 1, 0])))
 
 
+class gatherTests(unittest.TestCase):
+    def testGatherNumerical1(self):
+        t = TensorBase(np.array([[65, 17], [14, 25], [76, 22]]))
+        idx = TensorBase(np.array([[0], [1], [0]]))
+        dim = 1
+        result = t.gather(dim=dim, index=idx)
+        self.assertTrue(np.array_equal(result.data, np.array([[65], [25], [76]])))
+
+    def testGatherNumerical2(self):
+        t = TensorBase(np.array([[47, 74, 44], [56, 9, 37]]))
+        idx = TensorBase(np.array([[0, 0, 1], [1, 1, 0], [0, 1, 0]]))
+        dim = 0
+        result = t.gather(dim=dim, index=idx)
+        expexted = [[47, 74, 37], [56, 9, 44.], [47, 9, 44]]
+        self.assertTrue(np.array_equal(result.data, np.array(expexted)))
+
+
 class scatterTests(unittest.TestCase):
     def testScatter_Numerical0(self):
         t = TensorBase(np.zeros((3, 5)))