array_split pr

python/mxnet/ndarray/numpy/_op.py

@@ -40,7 +40,7 @@
  'blackman', 'flip', 'around', 'hypot', 'bitwise_xor', 'bitwise_or', 'rad2deg', 'deg2rad', 'unique', 'lcm',
  'tril', 'identity', 'take', 'ldexp', 'vdot', 'inner', 'outer', 'equal', 'not_equal', 'greater', 'less',
  'greater_equal', 'less_equal', 'hsplit', 'rot90', 'einsum', 'true_divide', 'nonzero', 'shares_memory',
- 'may_share_memory', 'diff', 'resize', 'nan_to_num', 'where']
+ 'may_share_memory', 'diff', 'resize', 'nan_to_num', 'where', 'array_split']
 
 
 @set_module('mxnet.ndarray.numpy')
@@ -3029,6 +3029,71 @@ def split(ary, indices_or_sections, axis=0):
 # pylint: enable=redefined-outer-name
 
 
+# pylint: disable=redefined-outer-name
+@set_module('mxnet.ndarray.numpy')
+def array_split(ary, indices_or_sections, axis=0):
+ """Split an array into multiple sub-arrays.
+
+ If `indices_or_sections` is an integer, N, the array will be divided
+ into N equal arrays along `axis`. If such a split is not possible,
+ an array of length l that should be split into n sections, it returns
+ l % n sub-arrays of size l//n + 1 and the rest of size l//n.
+
+ If `indices_or_sections` is a 1-D array of sorted integers, the entries
+ indicate where along `axis` the array is split. For example,
+ ``[2, 3]`` would, for ``axis=0``, result in
+ - ary[:2]
+ - ary[2:3]
+ - ary[3:]
+ If an index exceeds the dimension of the array along `axis`,
+ an empty sub-array is returned correspondingly.
+
+ Parameters
+ ----------
+ ary : ndarray
+ Array to be divided into sub-arrays.
+ indices_or_sections : int or 1-D Python tuple, list or set.
+ Param used to determine the number and size of the subarray.
+ axis : int, optional
+ The axis along which to split, default is 0.
+
+ Returns
+ -------
+ sub-arrays : list of ndarrays
+ A list of sub-arrays.
+
+ Examples
+ --------
+ >>> x = np.arange(9.0)
+ >>> np.array_split(x, 3)
+ [array([0., 1., 2.]), array([3., 4., 5.]), array([6., 7., 8.])]
+
+ >>> np.array_split(x, [3, 5, 6, 8])
+ [array([0., 1., 2.]), array([3., 4.]), array([5.]), array([6., 7.]), array([])]
+
+ >>> x = np.arange(8.0)
+ >>> np.array_split(x, 3)
+ [array([0., 1., 2.]), array([3., 4., 5.]), array([6., 7.])]
+
+ >>> x = np.arange(7.0)
+ >>> np.array_split(x, 3)
+ [array([0., 1., 2.]), array([3., 4.]), array([5., 6.])]
+ """
+ indices = []
+ sections = 0
+ if isinstance(indices_or_sections, integer_types):
+ sections = indices_or_sections
+ elif isinstance(indices_or_sections, (list, set, tuple)):
+ indices = [0] + list(indices_or_sections)
+ else:
+ raise ValueError('indices_or_sections must be either int, or tuple / list / set of ints')
+ ret = _npi.array_split(ary, indices, axis, False, sections)
+ if not isinstance(ret, list):
+ return [ret]
+ return ret
+# pylint: enable=redefined-outer-name
+
+
 # pylint: disable=redefined-outer-name
 @set_module('mxnet.ndarray.numpy')
 def hsplit(ary, indices_or_sections):

python/mxnet/numpy/multiarray.py

@@ -59,7 +59,7 @@
  'bitwise_xor', 'bitwise_or', 'rad2deg', 'deg2rad', 'unique', 'lcm', 'tril', 'identity', 'take',
  'ldexp', 'vdot', 'inner', 'outer', 'equal', 'not_equal', 'greater', 'less', 'greater_equal',
  'less_equal', 'hsplit', 'rot90', 'einsum', 'true_divide', 'nonzero', 'shares_memory',
- 'may_share_memory', 'diff', 'resize', 'nan_to_num', 'where']
+ 'may_share_memory', 'diff', 'resize', 'nan_to_num', 'where', 'array_split']
 
 # Return code for dispatching indexing function call
 _NDARRAY_UNSUPPORTED_INDEXING = -1
@@ -4841,6 +4841,58 @@ def split(ary, indices_or_sections, axis=0):
  return _mx_nd_np.split(ary, indices_or_sections, axis=axis)
 
 
+@set_module('mxnet.numpy')
+def array_split(ary, indices_or_sections, axis=0):
+ """Split an array into multiple sub-arrays.
+
+ If `indices_or_sections` is an integer, N, the array will be divided
+ into N equal arrays along `axis`. If such a split is not possible,
+ an array of length l that should be split into n sections, it returns
+ l % n sub-arrays of size l//n + 1 and the rest of size l//n.
+
+ If `indices_or_sections` is a 1-D array of sorted integers, the entries
+ indicate where along `axis` the array is split. For example,
+ ``[2, 3]`` would, for ``axis=0``, result in
+ - ary[:2]
+ - ary[2:3]
+ - ary[3:]
+ If an index exceeds the dimension of the array along `axis`,
+ an empty sub-array is returned correspondingly.
+
+ Parameters
+ ----------
+ ary : ndarray
+ Array to be divided into sub-arrays.
+ indices_or_sections : int or 1-D Python tuple, list or set.
+ Param used to determine the number and size of the subarray.
+ axis : int, optional
+ The axis along which to split, default is 0.
+
+ Returns
+ -------
+ sub-arrays : list of ndarrays
+ A list of sub-arrays.
+
+ Examples
+ --------
+ >>> x = np.arange(9.0)
+ >>> np.array_split(x, 3)
+ [array([0., 1., 2.]), array([3., 4., 5.]), array([6., 7., 8.])]
+
+ >>> np.array_split(x, [3, 5, 6, 8])
+ [array([0., 1., 2.]), array([3., 4.]), array([5.]), array([6., 7.]), array([])]
+
+ >>> x = np.arange(8.0)
+ >>> np.array_split(x, 3)
+ [array([0., 1., 2.]), array([3., 4., 5.]), array([6., 7.])]
+
+ >>> x = np.arange(7.0)
+ >>> np.array_split(x, 3)
+ [array([0., 1., 2.]), array([3., 4.]), array([5., 6.])]
+ """
+ return _mx_nd_np.array_split(ary, indices_or_sections, axis=axis)
+
+
 @set_module('mxnet.numpy')
 def vsplit(ary, indices_or_sections):
  r"""

python/mxnet/numpy_dispatch_protocol.py

@@ -113,6 +113,7 @@ def _run_with_array_ufunc_proto(*args, **kwargs):
  'reshape',
  'roll',
  'split',
+ 'array_split',
  'squeeze',
  'stack',
  'std',

python/mxnet/symbol/numpy/_symbol.py

@@ -40,7 +40,7 @@
  'add', 'subtract', 'multiply', 'divide', 'mod', 'remainder', 'power', 'arctan2',
  'sin', 'cos', 'tan', 'sinh', 'cosh', 'tanh', 'log10', 'sqrt', 'cbrt', 'abs', 'absolute', 'exp',
  'expm1', 'arcsin', 'arccos', 'arctan', 'sign', 'log', 'degrees', 'log2', 'log1p',
- 'rint', 'radians', 'reciprocal', 'square', 'negative', 'fix', 'ceil', 'floor',
+ 'rint', 'radians', 'reciprocal', 'square', 'negative', 'fix', 'ceil', 'floor', 'array_split',
  'trunc', 'logical_not', 'arcsinh', 'arccosh', 'arctanh', 'argsort', 'tensordot', 'histogram', 'eye',
  'linspace', 'logspace', 'expand_dims', 'tile', 'arange', 'split', 'vsplit', 'concatenate', 'append',
  'stack', 'vstack', 'column_stack', 'dstack', 'average', 'mean', 'maximum', 'minimum', 'swapaxes', 'clip',
@@ -3116,6 +3116,54 @@ def split(ary, indices_or_sections, axis=0):
 # pylint: enable=redefined-outer-name
 
 
+# pylint: disable=redefined-outer-name
+@set_module('mxnet.symbol.numpy')
+def array_split(ary, indices_or_sections, axis=0):
+ """Split an array into multiple sub-arrays.
+
+ If `indices_or_sections` is an integer, N, the array will be divided
+ into N equal arrays along `axis`. If such a split is not possible,
+ an array of length l that should be split into n sections, it returns
+ l % n sub-arrays of size l//n + 1 and the rest of size l//n.
+
+ If `indices_or_sections` is a 1-D array of sorted integers, the entries
+ indicate where along `axis` the array is split. For example,
+ ``[2, 3]`` would, for ``axis=0``, result in
+ - ary[:2]
+ - ary[2:3]
+ - ary[3:]
+ If an index exceeds the dimension of the array along `axis`,
+ an empty sub-array is returned correspondingly.
+
+ Parameters
+ ----------
+ ary : _Symbol
+ Array to be divided into sub-arrays.
+ indices_or_sections : int or 1-D Python tuple, list or set.
+ Param used to determine the number and size of the subarray.
+ axis : int, optional
+ The axis along which to split, default is 0.
+
+ Returns
+ -------
+ sub-arrays : list of ndarrays
+ A list of sub-arrays.
+ """
+ indices = []
+ sections = 0
+ if isinstance(indices_or_sections, int):
+ sections = indices_or_sections
+ elif isinstance(indices_or_sections, (list, set, tuple)):
+ indices = [0] + list(indices_or_sections)
+ else:
+ raise ValueError('indices_or_sections must either int or tuple / list / set of ints')
+ ret = _npi.array_split(ary, indices, axis, False, sections)
+ if not isinstance(ret, list):
+ return [ret]
+ return ret
+# pylint: enable=redefined-outer-name
+
+
 # pylint: disable=redefined-outer-name
 @set_module('mxnet.symbol.numpy')
 def hsplit(ary, indices_or_sections):

src/operator/tensor/matrix_op-inl.h

@@ -2727,9 +2727,15 @@ struct SplitParam : public dmlc::Parameter<SplitParam> {
 inline mxnet::TShape GetSplitIndices(const mxnet::TShape& ishape, int axis, int sections) {
  mxnet::TShape indices(sections+1, -1);
  indices[0] = 0;
- int64_t section_size = ishape[axis] / sections;
+ int64_t section_size_b = (int64_t) (ishape[axis] / sections);
+ int64_t section_size_a = section_size_b + 1;
+ int section_a = ishape[axis] % sections;
  for (int i = 0; i < sections; ++i) {
- indices[i+1] = section_size * (i + 1);
+ if( i < section_a){
+ indices[i+1] = section_size_a * (i + 1);
+ } else {
+ indices[i+1] = section_size_b + indices[i];
+ }
  }
  return indices;
 }

src/operator/tensor/matrix_op.cc

@@ -1038,6 +1038,7 @@ Example::
 
 NNVM_REGISTER_OP(_split_v2)
 .add_alias("_npi_split")
+.add_alias("_npi_array_split")
 .describe(R"code(Splits an array along a particular axis into multiple sub-arrays.
 Example::
  x = [[[ 1.]
@@ -1106,6 +1107,7 @@ Example::
 .add_arguments(SplitParam::__FIELDS__());
 
 NNVM_REGISTER_OP(_split_v2_backward)
+.add_alias("_backward_npi_array_split")
 .set_attr_parser(ParamParser<SplitParam>)
 .set_num_inputs(SplitNumOutputs)
 .set_num_outputs(1)

src/operator/tensor/matrix_op.cu

@@ -219,9 +219,11 @@ NNVM_REGISTER_OP(space_to_depth)
 .set_attr<FCompute>("FCompute<gpu>", SpaceToDepthOpForward<gpu>);
 
 NNVM_REGISTER_OP(_split_v2)
+.add_alias("_npi_array_split")
 .set_attr<FCompute>("FCompute<gpu>", SplitOpForward<gpu>);
 
 NNVM_REGISTER_OP(_split_v2_backward)
+.add_alias("_backward_npi_array_split")
 .set_attr<FCompute>("FCompute<gpu>", SplitOpBackward<gpu>);
 
 } // namespace op

tests/python/unittest/test_numpy_interoperability.py

@@ -181,6 +181,18 @@ def _add_workload_split():
  assertRaises(ValueError, np.split, np.arange(10), 3)
 
 
+def _add_workload_array_split():
+ a = np.arange(10)
+ b = np.array([np.arange(10), np.arange(10)])
+
+ for i in range(1, 12):
+ OpArgMngr.add_workload('array_split', a, i)
+ OpArgMngr.add_workload('array_split', b, 3, axis=0)
+ OpArgMngr.add_workload('array_split', b, [0, 1, 2], axis=0)
+ OpArgMngr.add_workload('array_split', b, 3, axis=-1)
+ OpArgMngr.add_workload('array_split', b, 3)
+
+
 def _add_workload_squeeze():
  OpArgMngr.add_workload('squeeze', np.random.uniform(size=(4, 1)))
  OpArgMngr.add_workload('squeeze', np.random.uniform(size=(20, 10, 10, 1, 1)))
@@ -1398,6 +1410,7 @@ def _prepare_workloads():
  _add_workload_rint(array_pool)
  _add_workload_roll()
  _add_workload_split()
+ _add_workload_array_split()
  _add_workload_squeeze()
  _add_workload_stack(array_pool)
  _add_workload_std()

tests/python/unittest/test_numpy_op.py

@@ -2237,6 +2237,63 @@ def get_indices(axis_size):
  assert_almost_equal(mx_out.asnumpy(), np_out, rtol=1e-3, atol=1e-5)
 
 
+@with_seed()
+@use_np
+def test_np_array_split():
+ class TestArray_split(HybridBlock):
+ def __init__(self, indices_or_sections, axis=None):
+ super(TestArray_split, self).__init__()
+ self._axis = axis
+ self._indices_or_sections = indices_or_sections
+
+ def hybrid_forward(self, F, a, *args, **kwargs):
+ return F.np.array_split(a, indices_or_sections=self._indices_or_sections,
+ axis=self._axis)
+
+ def get_indices(axis_size):
+ if axis_size is 0:
+ axis_size = random.randint(3, 6)
+ samples = random.randint(1, axis_size - 1)
+ indices = sorted(random.sample([i for i in range(0, axis_size + 1)], samples))
+ indices = tuple(indices)
+ return indices
+
+ shapes = [(), (5, ), (10, ),
+ (2, 5), (5, 5), (10, 10),
+ (4, 4, 4), (4, 6, 9), (6, 6, 6),
+ (7, 8, 9, 10)]
+ dtypes = [np.int8, np.uint8, np.int32, np.int64, np.float16, np.float32, np.float64]
+
+ combinations = itertools.product([False, True], shapes, dtypes)
+ for hybridize, shape, dtype in combinations:
+ rtol = 1e-2 if dtype == np.float16 else 1e-3
+ atol = 1e-4 if dtype == np.float16 else 1e-5
+ for axis in range(len(shape)):
+ x = np.random.uniform(-5.0, 5.0, size=shape).astype(dtype)
+ indices = get_indices(shape[axis])
+ sections = 7 if x.shape[axis] is 0 else random.randint(1,x.shape[axis])
+ for indices_or_sections in [indices, sections]:
+ # test gluon
+ test_array_split = TestArray_split(axis=axis, indices_or_sections=indices_or_sections)
+ if hybridize:
+ test_array_split.hybridize()
+ x.attach_grad()
+ expected_ret = _np.array_split(x.asnumpy(), indices_or_sections=indices_or_sections, axis=axis)
+ with mx.autograd.record():
+ y = test_array_split(x)
+ assert len(y) == len(expected_ret)
+ for mx_out, np_out in zip(y, expected_ret):
+ assert_almost_equal(mx_out.asnumpy(), np_out, rtol=rtol, atol=atol)
+ mx.autograd.backward(y)
+ assert_almost_equal(x.grad.asnumpy(), _np.ones(x.shape), rtol=rtol, atol=atol)
+
+ # test imperative
+ mx_outs = np.array_split(x, indices_or_sections=indices_or_sections, axis=axis)
+ np_outs = _np.array_split(x.asnumpy(), indices_or_sections=indices_or_sections, axis=axis)
+ for mx_out, np_out in zip(mx_outs, np_outs):
+ assert_almost_equal(mx_out.asnumpy(), np_out, rtol=rtol, atol=atol)
+
+
 @with_seed()
 @use_np
 def test_np_vsplit():