MAINT: move *split logic to _impl

Author: ev-br

torch_np/_detail/implementations.py

@@ -155,6 +155,8 @@ def split_helper_int(tensor, indices_or_sections, axis, strict=False):
     if not isinstance(indices_or_sections, int):
         raise NotImplementedError("split: indices_or_sections")
 
+    axis = _util.normalize_axis_index(axis, tensor.ndim)
+
     # numpy: l%n chunks of size (l//n + 1), the rest are sized l//n
     l, n = tensor.shape[axis], indices_or_sections
 
@@ -195,6 +197,26 @@ def split_helper_list(tensor, indices_or_sections, axis):
     return torch.split(tensor, lst, axis)
 
 
+def hsplit(tensor, indices_or_sections):
+    if tensor.ndim == 0:
+        raise ValueError("hsplit only works on arrays of 1 or more dimensions")
+    axis = 1 if tensor.ndim > 1 else 0
+    return split_helper(tensor, indices_or_sections, axis, strict=True)
+
+
+def vsplit(tensor, indices_or_sections):
+    if tensor.ndim < 2:
+        raise ValueError("vsplit only works on arrays of 2 or more dimensions")
+    return split_helper(tensor, indices_or_sections, 0, strict=True)
+
+
+def dsplit(tensor, indices_or_sections):
+    if tensor.ndim < 3:
+        raise ValueError("dsplit only works on arrays of 3 or more dimensions")
+    return split_helper(tensor, indices_or_sections, 2, strict=True)
+
+
+
 def clip(tensor, t_min, t_max):
     if t_min is not None:
         t_min = torch.broadcast_to(t_min, tensor.shape)

torch_np/_wrapper.py

@@ -176,56 +176,35 @@ 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()
     base = ary if isinstance(ary, ndarray) else None
-    axis = _util.normalize_axis_index(axis, tensor.ndim)
-
     result = _impl.split_helper(tensor, indices_or_sections, axis)
-
     return tuple(maybe_set_base(x, base) for x in result)
 
 
 def split(ary, indices_or_sections, axis=0):
     tensor = asarray(ary).get()
     base = ary if isinstance(ary, ndarray) else None
-    axis = _util.normalize_axis_index(axis, tensor.ndim)
-
     result = _impl.split_helper(tensor, indices_or_sections, axis, strict=True)
-
     return tuple(maybe_set_base(x, base) for x in result)
 
 
 def hsplit(ary, indices_or_sections):
     tensor = asarray(ary).get()
     base = ary if isinstance(ary, ndarray) else None
-
-    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 = _impl.split_helper(tensor, indices_or_sections, axis, strict=True)
-
+    result = _impl.hsplit(tensor, indices_or_sections)
     return tuple(maybe_set_base(x, base) for x in result)
 
 
 def vsplit(ary, indices_or_sections):
     tensor = asarray(ary).get()
     base = ary if isinstance(ary, ndarray) else None
-
-    if tensor.ndim < 2:
-        raise ValueError("vsplit only works on arrays of 2 or more dimensions")
-    result = _impl.split_helper(tensor, indices_or_sections, 0, strict=True)
-
+    result = _impl.vsplit(tensor, indices_or_sections)
     return tuple(maybe_set_base(x, base) for x in result)
 
 
 def dsplit(ary, indices_or_sections):
     tensor = asarray(ary).get()
     base = ary if isinstance(ary, ndarray) else None
-
-    if tensor.ndim < 3:
-        raise ValueError("dsplit only works on arrays of 3 or more dimensions")
-    result = _impl.split_helper(tensor, indices_or_sections, 2, strict=True)
-
+    result = _impl.dsplit(tensor, indices_or_sections)
     return tuple(maybe_set_base(x, base) for x in result)