WIP: splits

Author: ev-br

torch_np/_wrapper.py

@@ -158,6 +158,88 @@ def stack(arrays, axis=0, out=None, *, dtype=None, casting="same_kind"):
     )
 
 
+def array_split(ary, indices_or_sections, axis=0):
+    tensor = asarray(ary).get()
+    axis = _util.normalize_axis_index(axis, tensor.ndim)
+
+    result = _split_helper(tensor, indices_or_sections, axis)
+
+    return tuple(asarray(_) for _ in result)
+
+
+def split(ary, indices_or_sections, axis=0):
+    tensor = asarray(ary).get()
+    axis = _util.normalize_axis_index(axis, tensor.ndim)
+
+    result = _split_helper(tensor, indices_or_sections, axis, strict=True)
+
+    return tuple(asarray(_) for _ in result)
+
+
+def hsplit(ary, indices_or_sections):
+    tensor = asarray(ary).get()
+
+    if tensor.ndim == 0:
+        raise ValueError('hsplit only works on arrays of 1 or more dimensions')
+
+    axis = 1 if tensor.ndim > 1 else 0
+
+    result = _split_helper(tensor, indices_or_sections, axis, strict=True)
+
+    return tuple(asarray(_) for _ in result)
+
+
+def vsplit(ary, indices_or_sections):
+    tensor = asarray(ary).get()
+
+    if tensor.ndim < 2:
+        raise ValueError('vsplit only works on arrays of 2 or more dimensions')
+    result = _split_helper(tensor, indices_or_sections, 0, strict=True)
+
+    return tuple(asarray(_) for _ in result)
+
+
+def dsplit(ary, indices_or_sections):
+    tensor = asarray(ary).get()
+
+    if tensor.ndim < 3:
+        raise ValueError('dsplit only works on arrays of 3 or more dimensions')
+    result = _split_helper(tensor, indices_or_sections, 2, strict=True)
+
+    return tuple(asarray(_) for _ in result)
+
+
+def _split_helper(tensor, indices_or_sections, axis, strict=False):
+    if not isinstance(indices_or_sections, int):
+        raise NotImplementedError('split: indices_or_sections')
+
+    # numpy: l%n chunks of size (l//n + 1), the rest are sized l//n
+    l, n = tensor.shape[axis], indices_or_sections
+
+    if n <= 0:
+        raise ValueError()
+
+    if l % n == 0:
+        num, sz = n, l // n
+        lst = [sz] * num
+    else:
+        if strict:
+            raise ValueError("array split does not result in an equal division")
+
+        num, sz = l % n, l // n + 1
+        lst = [sz] * num
+
+        lrest = l - num*sz
+
+        sz_1 = sz - 1
+        num_1 = lrest // sz_1
+        lst += [sz_1]*num_1
+
+    result = torch.split(tensor, lst, axis)
+
+    return result
+
+
 def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None, axis=0):
     if axis != 0 or retstep or not endpoint:
         raise NotImplementedError

torch_np/tests/numpy_tests/lib/test_shape_base_.py

@@ -2,12 +2,12 @@
 import sys
 import pytest
 
-from numpy.lib.shape_base import (apply_along_axis, apply_over_axes, array_split,
-    split, hsplit, dsplit, vsplit, kron, tile,
-    expand_dims, take_along_axis, put_along_axis)
+from numpy.lib.shape_base import (apply_along_axis, apply_over_axes, kron, tile,
+    take_along_axis, put_along_axis)
 
 import torch_np as np
-from torch_np import column_stack, dstack, expand_dims
+from torch_np import (column_stack, dstack, expand_dims, array_split,
+    split, hsplit, dsplit, vsplit,)
 
 from torch_np.random import rand
 
@@ -275,7 +275,6 @@ def test_repeated_axis(self):
         assert_raises(ValueError, expand_dims, a, axis=(1, 1))
 
 
-@pytest.mark.xfail(reason="TODO: implement")
 class TestArraySplit:
     def test_integer_0_split(self):
         a = np.arange(10)
@@ -410,7 +409,6 @@ def test_index_split_high_bound(self):
         compare_results(res, desired)
 
 
-@pytest.mark.xfail(reason="TODO: implement")
 class TestSplit:
     # The split function is essentially the same as array_split,
     # except that it test if splitting will result in an
@@ -493,7 +491,6 @@ def test_generator(self):
 
 # array_split has more comprehensive test of splitting.
 # only do simple test on hsplit, vsplit, and dsplit
-@pytest.mark.xfail(reason="TODO: implement")
 class TestHsplit:
     """Only testing for integer splits.
 
@@ -523,7 +520,6 @@ def test_2D_array(self):
         compare_results(res, desired)
 
 
-@pytest.mark.xfail(reason="TODO: implement")
 class TestVsplit:
     """Only testing for integer splits.
 
@@ -551,7 +547,6 @@ def test_2D_array(self):
         compare_results(res, desired)
 
 
-@pytest.mark.xfail(reason="TODO: implement")
 class TestDsplit:
     # Only testing for integer splits.
     def test_non_iterable(self):