Repository: OpenMined/PySyft (python)

Issue:
Implement Default split Functionality for Base Tensor Type.
<!-- Please make sure that you review this: https://github.com/OpenMined/Docs/blob/master/contributing/guidelines.md  -->
<!-- If you are looking to file a bug make sure to look at the .github/BUG_ISSUE_TEMPLATE.md -->


#### User story: 
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 `split` should return a new tensor. For a reference on the operation these perform check out [PyTorch's](http://pytorch.org/docs/master/torch.html#torch.split) documentation.

<!-- Provide a detailed explaination about the proposed feature, you can draw inspiration from something like this: 
https://github.com/OpenMined/PySyft/issues/227 or https://github.com/OpenMined/PySyft/issues/12 -->

#### Acceptance Criteria: 

<!-- Provide an outline af all the things that needs to be addressed in order to close this Issue,
be as descriptive as possible -->
  - [x] If the Base Tensor type's attribute "encrypted" is set to True, it should return a NotImplemented error.
- [ ] corresponding unit tests demonstrating the correct operation on the Base Tensor type implemented over int and float Tensors.
 - [ ] inline documentation as described over [here](https://github.com/OpenMined/PySyft/blob/85bc68e81a2f4bfc0f0bf6c4252b88d6d7b54004/syft/math.py#L5). 


<!-- Thanks for your contributions! -->


Pull request description:
Split tensor
This fixes #383 

#### What does your implementation fix? 
This add split functionality to default tensor type.

#### Explain your changes.
In syft.math it add a split function which uses numpy.array_split() for splitting the array into given axis into `split size` size of each chunk. The function returns a list containing divided sub-arrays. I had to some initial calculations to match the numpy api. 

Patch:
diff --git a/syft/math.py b/syft/math.py
index e880c85ed9..4d395a6a0d 100644
--- a/syft/math.py
+++ b/syft/math.py
@@ -968,3 +968,43 @@ def zeros(dim):
         Output Tensor
     """
     return TensorBase(np.zeros(dim))
+
+
+def split(tensor, split_size, axis=0):
+    """
+    Splits the tensor into multiple equally sized chunks (if possible).
+
+    Last chunk will be smaller if the tensor size along a given axis
+    is not divisible by `split_size`.
+
+    Returns a list of the split tensors
+
+    Parameters
+    ----------
+    tensor: TensorBase
+        array to be divided into sub-arrays.
+
+    split_size: int
+        size of single chunk
+
+    axis: int, optional
+        The axis along which to split, default is 0.
+
+    Returns
+    -------
+    list: list of divided arrays
+    """
+    if axis < 0:
+        axis += len(tensor.shape())
+
+    length_along_axis = tensor.shape()[axis]
+
+    # calculate number of splits!
+    num_splits = (length_along_axis + split_size - 1) // split_size
+
+    # make array to pass to numpy array_split function
+    split_according = [split_size * i for i in range(1, num_splits)]
+
+    list_ = np.array_split(tensor.data, split_according, axis)
+
+    return list(map(lambda x: TensorBase(x), list_))
diff --git a/syft/tensor.py b/syft/tensor.py
index f774b294e1..89f23c3b05 100644
--- a/syft/tensor.py
+++ b/syft/tensor.py
@@ -3038,29 +3038,6 @@ class TensorBase(object):
         else:
             return self.data.shape
 
-    def split(self, split_size, dim=0):
-        """
-        Returns tuple of tensors of equally sized tensor/chunks (if possible)
-
-        Parameters
-        ----------
-        split_size:
-
-        dim:
-
-        Returns
-        -------
-        Tuple of Tensors
-        """
-        if self.encrypted:
-            return NotImplemented
-        splits = np.array_split(self.data, split_size, axis=0)
-        tensors = list()
-        for s in splits:
-            tensors.append(TensorBase(s))
-        tensors_tuple = tuple(tensors)
-        return tensors_tuple
-
     def stride(self, dim=None):
         """
         Returns the jump necessary to go from one element to the next one in the specified dimension dim.
@@ -3601,3 +3578,30 @@ class TensorBase(object):
 
         self.data.fill(0)
         return self
+
+    def split(self, split_size, axis=0):
+        """
+        Splits the tensor into multiple equally sized chunks (if possible).
+
+        Last chunk will be smaller if the tensor size along a given axis
+        is not divisible by `split_size`.
+
+        Returns a list of the split tensors
+
+        Parameters
+        ----------
+
+        split_size: int
+            size of single chunk
+
+        axis: int, optional
+            The axis along which to split, default is 0.
+
+        Returns
+        -------
+        list: list of divided arrays
+        """
+        if self.encrypted:
+            return NotImplemented
+
+        return syft.math.split(self, split_size, axis)


Test patch:
diff --git a/tests/test_math.py b/tests/test_math.py
index 83b590f47b..09aad0a9f7 100644
--- a/tests/test_math.py
+++ b/tests/test_math.py
@@ -427,3 +427,17 @@ class MultinomialTests(unittest.TestCase):
         t2 = syft.math.multinomial(t1, len(t1))
         self.assertTupleEqual((len(t1),), t2.shape())
         self.assertTrue(np.all(t2.data >= 0) and np.all(t2.data <= len(t1)))
+
+
+class SplitTests(unittest.TestCase):
+    def test_split(self):
+        t = TensorBase(np.random.rand(7, 4))
+        split_size = 3
+        axis = 0
+        target_shapes = [(3, 4), (3, 4), (1, 4)]
+        splits = syft.math.split(t, split_size, axis)
+        start = 0
+        for target_shape, split in zip(target_shapes, splits):
+            self.assertTrue(syft.equal(split.shape(), target_shape))
+            self.assertTrue(syft.equal(t.narrow(axis, start, target_shape[axis]), split))
+            start += target_shape[axis]
diff --git a/tests/test_tensor.py b/tests/test_tensor.py
index c34b4cebc6..3c9052c3bf 100644
--- a/tests/test_tensor.py
+++ b/tests/test_tensor.py
@@ -974,13 +974,6 @@ class CumprodTest(unittest.TestCase):
         self.assertTrue(np.equal(t1.cumprod_(dim=1), t3).all())
 
 
-class SplitTests(unittest.TestCase):
-    def test_split(self):
-        t1 = TensorBase(np.arange(8.0))
-        t2 = t1.split(4)
-        self.assertTrue(np.array_equal(t2, tuple((np.array([0., 1.]), np.array([2., 3.]), np.array([4., 5.]), np.array([6., 7.])))))
-
-
 class SqueezeTests(unittest.TestCase):
     def test_squeeze(self):
         t1 = TensorBase(np.zeros((2, 1, 2, 1, 2)))
@@ -1628,6 +1621,20 @@ class UnfoldTest(unittest.TestCase):
                                        t1_unfolded_actual_2))
 
 
+class SplitTests(unittest.TestCase):
+    def test_split(self):
+        t = TensorBase(np.random.rand(10, 5))
+        split_size = 3
+        axis = 0
+        target_shapes = [(3, 5), (3, 5), (3, 5), (1, 5)]
+        splits = t.split(split_size, axis)
+        start = 0
+        for target_shape, split in zip(target_shapes, splits):
+            self.assertTrue(syft.equal(split.shape(), target_shape))
+            self.assertTrue(syft.equal(t.narrow(axis, start, target_shape[axis]), split))
+            start += target_shape[axis]
+
+
 if __name__ == "__main__":
 
     unittest.main()