Repository: OpenMined/PySyft

Issue:
Implement Default index 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, index() and index_select() should return a new tensor and index_add_(), index_copy_(), and index_fill_() should perform the operation 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:
samuxiii: Ok, I've added slices into the index as a possible type of parameter. But.. I don't understand the comment about index_add_, I've almost used the same pytorch example as my unit test.
Please, could you be more explicit? Thanks
bharathgs: @samuxiii  ahh, my bad.did not see the `index_add_` properly. It looks good.

Patch:
diff --git a/syft/tensor.py b/syft/tensor.py
index f5313aa491..719ae2d409 100644
--- a/syft/tensor.py
+++ b/syft/tensor.py
@@ -2848,6 +2848,97 @@ class TensorBase(object):
         self.data = np.remainder(self.data, divisor)
         return self
 
+    def index(self, m):
+        """
+        Returns a new Tensor with the element selected by position
+
+        :param m: integer index or slice
+        :return: tensor of selected indices
+        """
+        if self.encrypted:
+            return NotImplemented
+        if not isinstance(m, int) and not isinstance(m, slice):
+            raise ValueError("The value of index must be integer")
+        return TensorBase(self.data[m], self.encrypted)
+
+    def index_add_(self, dim, index, tensor):
+        """
+        Add the value of 'tensor' selecting the elements and ordered
+        by index. In-place operation.
+
+        :param dim: dimension along which to index
+        :param index: 1D tensor containing the indices to select
+        :param tensor: tensor containing the values to add
+        """
+        index = _ensure_tensorbase(index)
+        tensor = _ensure_tensorbase(tensor)
+
+        if self.encrypted:
+            return NotImplemented
+        if index.data.dtype != np.dtype('int_'):
+            raise TypeError("The value of index must be integer")
+        if self.data.shape != tensor.data.shape:
+            raise IndexError("Tensor has different shape")
+        if self.data.shape[dim] != index.data.size:
+            raise ValueError("Index should have the same number of elements as dimension")
+        if np.argmax(index.data > self.data.shape[dim]) != 0:
+            raise ValueError("Index contains a value which is out of range")
+        if dim >= self.data.ndim or dim < -self.data.ndim:
+            raise IndexError("Dimension out of range")
+
+        self.data += tensor.data.take(index, dim)
+
+    def index_copy_(self, dim, index, tensor):
+        """
+        Copy the values of 'tensor' selecting the elements and ordered
+        by index. In-place operation.
+
+        :para dim: dimension along which to index
+        :param index: 1D tensor containing the indices to select
+        :param tensor: tensor containing the values to add
+        """
+        index = _ensure_tensorbase(index)
+        tensor = _ensure_tensorbase(tensor)
+
+        if self.encrypted:
+            return NotImplemented
+        if index.data.dtype != np.dtype('int_'):
+            raise TypeError("The value of index must be integer")
+        if self.data.shape != tensor.data.shape:
+            raise IndexError("Tensor has different shape")
+        if self.data.shape[dim] != index.data.size:
+            raise ValueError("Index should have the same number of elements as dimension")
+        if np.argmax(index.data > self.data.shape[dim]) != 0:
+            raise ValueError("Index contains a value which is out of range")
+        if dim >= self.data.ndim or dim < -self.data.ndim:
+            raise IndexError("Dimension out of range")
+
+        np.copyto(self.data, tensor.data.take(index, dim))
+
+    def index_fill_(self, dim, index, value):
+        """
+        Fill the original tensor with the values of 'tensor' selecting
+        the elements and ordered by index. In-place operation.
+
+        :param dim: dimension along which to inde
+        :param index: 1D tensor containing the indices to select
+        :param value: value to fill
+        """
+        index = _ensure_tensorbase(index)
+
+        if self.encrypted:
+            return NotImplemented
+        if index.data.dtype != np.dtype('int_'):
+            raise TypeError("The value of index must be integer")
+        if np.argmax(index.data > self.data.shape[dim]) != 0:
+            raise ValueError("Index contains a value which is out of range")
+        if dim >= self.data.ndim or dim < -self.data.ndim:
+            raise IndexError("Dimension out of range")
+
+        idx = [slice(None)] * self.data.ndim
+        idx[dim] = index
+        self.data[tuple(idx)] = value
+
     def index_select(self, dim, index):
         """
         Returns a new Tensor which indexes the ``input`` Tensor along


Test patch:
diff --git a/tests/test_tensor.py b/tests/test_tensor.py
index 465370d309..bf91ebdc3a 100644
--- a/tests/test_tensor.py
+++ b/tests/test_tensor.py
@@ -4,6 +4,7 @@ import unittest
 from syft import tensor
 import numpy as np
 import math
+import pytest
 
 
 # Here's our "unit tests".
@@ -250,13 +251,6 @@ class EqualTests(unittest.TestCase):
         self.assertTrue(t1 != t2)
 
 
-class IndexTests(unittest.TestCase):
-    def test_indexing(self):
-        t1 = TensorBase(np.array([1.2, 2, 3]))
-        self.assertEqual(1.2, t1[0])
-        self.assertEqual(3, t1[-1])
-
-
 class sigmoidTests(unittest.TestCase):
     def test_sigmoid(self):
         t1 = TensorBase(np.array([1.2, 3.3, 4]))
@@ -954,7 +948,90 @@ class notEqualTests(unittest.TestCase):
         self.assertTrue(syft.equal(t1, TensorBase([1, 1, 1, 0])))
 
 
-class index_selectTests(unittest.TestCase):
+class IndexTests(unittest.TestCase):
+    def test_indexing(self):
+        t1 = TensorBase(np.array([1.2, 2, 3]))
+        self.assertEqual(1.2, t1[0])
+        self.assertEqual(3, t1[-1])
+
+    def test_index(self):
+        t = TensorBase(np.array([1, 2, 3.5, 4, 5, 6, 3.5]))
+        expected1 = TensorBase(np.array(2))
+        expected2 = TensorBase(np.array(3.5))
+        expected3 = TensorBase(np.array([4, 5, 6]))
+
+        self.assertEqual(expected1, t.index(1))
+        self.assertEqual(expected2, t.index(2))
+        self.assertEqual(expected2, t.index(-1))
+        self.assertEqual(expected3, t.index(slice(3, 6)))
+        with pytest.raises(ValueError):
+            t.index(3.5)
+
+    def test_index_add_(self):
+        t1 = TensorBase(np.array([[0, 0, 0], [1, 1, 1], [1, 1, 1]]))
+        t2 = TensorBase(np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]))
+
+        expected_0 = TensorBase(np.array([[1, 2, 3], [8, 9, 10], [5, 6, 7]]))
+        t1.index_add_(0, [0, 2, 1], t2)
+        self.assertEqual(expected_0, t1)
+
+        t1 = TensorBase(np.array([[0, 0, 0], [1, 1, 1], [1, 1, 1]]))
+        expected_1 = TensorBase(np.array([[1, 3, 2], [5, 7, 6], [8, 10, 9]]))
+        t1.index_add_(1, [0, 2, 1], t2)
+        self.assertEqual(expected_1, t1)
+
+        with pytest.raises(TypeError):
+            t1.index_add_(0, [1.0, 2, 2], t2)
+        with pytest.raises(IndexError):
+            t1.index_add_(0, [0, 1, 2], TensorBase([1, 2]))
+        with pytest.raises(ValueError):
+            t1.index_add_(0, [0, 1], t2)
+        with pytest.raises(ValueError):
+            t1.index_add_(0, [0, 1, 5], t2)
+        with pytest.raises(IndexError):
+            t1.index_add_(4, [0, 1, 2], t2)
+
+    def test_index_copy_(self):
+        t1 = TensorBase(np.array([[0, 0, 0], [1, 1, 1], [1, 1, 1]]))
+        t2 = TensorBase(np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]))
+        expected_0 = TensorBase(np.array([[1, 2, 3], [7, 8, 9], [4, 5, 6]]))
+        t1.index_copy_(0, [0, 2, 1], t2)
+        self.assertEqual(expected_0, t1)
+
+        t1 = TensorBase(np.array([[0, 0, 0], [1, 1, 1], [1, 1, 1]]))
+        expected_1 = TensorBase(np.array([[3, 1, 2], [6, 4, 5], [9, 7, 8]]))
+        t1.index_copy_(1, [2, 0, 1], t2)
+        self.assertEqual(expected_1, t1)
+
+        with pytest.raises(TypeError):
+            t1.index_copy_(0, [1.0, 2, 2], t2)
+        with pytest.raises(IndexError):
+            t1.index_copy_(0, [0, 1, 2], TensorBase([1, 2]))
+        with pytest.raises(ValueError):
+            t1.index_copy_(0, [0, 1], t2)
+        with pytest.raises(ValueError):
+            t1.index_copy_(0, [0, 1, 5], t2)
+        with pytest.raises(IndexError):
+            t1.index_copy_(4, [0, 1, 2], t2)
+
+    def test_index_fill_(self):
+        t1 = TensorBase(np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]))
+        expected_0 = TensorBase(np.array([[1, 1, 1], [1, 1, 1], [7, 8, 9]]))
+        t1.index_fill_(0, [0, 1], 1)
+        self.assertEqual(expected_0, t1)
+
+        t1 = TensorBase(np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]))
+        expected_1 = TensorBase(np.array([[-2, 2, -2], [-2, 5, -2], [-2, 8, -2]]))
+        t1.index_fill_(1, [0, 2], -2)
+        self.assertEqual(expected_1, t1)
+
+        with pytest.raises(TypeError):
+            t1.index_fill_(0, [1.0, 2, 2], 1)
+        with pytest.raises(ValueError):
+            t1.index_fill_(0, [0, 1, 5], 1)
+        with pytest.raises(IndexError):
+            t1.index_fill_(4, [0, 1, 2], 1)
+
     def test_index_select(self):
         t = TensorBase(np.reshape(np.arange(0, 2 * 3 * 4), (2, 3, 4)))
         idx = np.array([1, 0])