Skip to content
This repository has been archived by the owner on Nov 17, 2023. It is now read-only.

Commit

Permalink
add numpy op indices
Browse files Browse the repository at this point in the history
  • Loading branch information
gyshi committed Aug 20, 2019
1 parent bcbdc1c commit b37a698
Show file tree
Hide file tree
Showing 7 changed files with 394 additions and 4 deletions.
73 changes: 71 additions & 2 deletions python/mxnet/ndarray/numpy/_op.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@
from ..ndarray import NDArray

__all__ = ['zeros', 'ones', 'add', 'subtract', 'multiply', 'divide', 'mod', 'power', 'tensordot',
'linspace', 'expand_dims', 'tile', 'arange', 'split', 'concatenate', 'stack']
'linspace', 'expand_dims', 'tile', 'arange', 'split', 'concatenate', 'stack', 'indices']


@set_module('mxnet.ndarray.numpy')
Expand Down Expand Up @@ -634,7 +634,7 @@ def tile(A, reps):
"""
return _unary_func_helper(A, _npi.tile, _np.tile, reps=reps)


# pylint: disable=redefined-outer-name
@set_module('mxnet.ndarray.numpy')
def split(ary, indices_or_sections, axis=0):
"""Split an array into multiple sub-arrays.
Expand Down Expand Up @@ -682,6 +682,7 @@ def split(ary, indices_or_sections, axis=0):
if not isinstance(ret, list):
return [ret]
return ret
# pylint: enable=redefined-outer-name


@set_module('mxnet.ndarray.numpy')
Expand Down Expand Up @@ -732,3 +733,71 @@ def get_list(arrays):

arrays = get_list(arrays)
return _npi.stack(*arrays, axis=axis, out=out)


# pylint: disable=redefined-outer-name
@set_module('mxnet.ndarray.numpy')
def indices(dimensions, dtype=_np.int32, ctx=None):
"""Return an array representing the indices of a grid.
Compute an array where the subarrays contain index values 0,1,...
varying only along the corresponding axis.
Parameters
----------
dimensions : sequence of ints
The shape of the grid.
dtype : data-type, optional
The desired data-type for the array. Default is `float32`.
ctx : device context, optional
Device context on which the memory is allocated. Default is
`mxnet.context.current_context()`.
Returns
-------
grid : ndarray
The array of grid indices,
``grid.shape = (len(dimensions),) + tuple(dimensions)``.
Notes
-----
The output shape is obtained by prepending the number of dimensions
in front of the tuple of dimensions, i.e. if `dimensions` is a tuple
``(r0, ..., rN-1)`` of length ``N``, the output shape is
``(N,r0,...,rN-1)``.
The subarrays ``grid[k]`` contains the N-D array of indices along the
``k-th`` axis. Explicitly::
grid[k,i0,i1,...,iN-1] = ik
Examples
--------
>>> grid = np.indices((2, 3))
>>> grid.shape
(2, 2, 3)
>>> grid[0] # row indices
array([[0, 0, 0],
[1, 1, 1]])
>>> grid[1] # column indices
array([[0, 0, 0],
[1, 1, 1]], dtype=int32)
The indices can be used as an index into an array.
>>> x = np.arange(20).reshape(5, 4)
>>> row, col = np.indices((2, 3))
>>> x[row, col]
array([[0., 1., 2.],
[4., 5., 6.]])
Note that it would be more straightforward in the above example to
extract the required elements directly with ``x[:2, :3]``.
"""
if isinstance(dimensions, (tuple, list)):
if ctx is None:
ctx = current_context()
return _npi.indices(dimensions=dimensions, dtype=dtype, ctx=ctx)
else:
raise ValueError("The dimensions must be sequence of ints")
# pylint: enable=redefined-outer-name
63 changes: 62 additions & 1 deletion python/mxnet/numpy/multiarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@

__all__ = ['ndarray', 'empty', 'array', 'zeros', 'ones', 'add', 'subtract', 'multiply', 'divide',
'mod', 'power', 'tensordot', 'linspace', 'expand_dims', 'tile', 'arange', 'split',
'concatenate', 'stack']
'concatenate', 'stack', 'indices']


# This function is copied from ndarray.py since pylint
Expand Down Expand Up @@ -1898,3 +1898,64 @@ def stack(arrays, axis=0, out=None):
stacked : ndarray
The stacked array has one more dimension than the input arrays."""
return _mx_nd_np.stack(arrays, axis=axis, out=out)


@set_module('mxnet.numpy')
def indices(dimensions, dtype=_np.int32, ctx=None):
"""Return an array representing the indices of a grid.
Compute an array where the subarrays contain index values 0,1,...
varying only along the corresponding axis.
Parameters
----------
dimensions : sequence of ints
The shape of the grid.
dtype : data-type, optional
The desired data-type for the array. Default is `float32`.
ctx : device context, optional
Device context on which the memory is allocated. Default is
`mxnet.context.current_context()`.
Returns
-------
grid : ndarray
The array of grid indices,
``grid.shape = (len(dimensions),) + tuple(dimensions)``.
Notes
-----
The output shape is obtained by prepending the number of dimensions
in front of the tuple of dimensions, i.e. if `dimensions` is a tuple
``(r0, ..., rN-1)`` of length ``N``, the output shape is
``(N,r0,...,rN-1)``.
The subarrays ``grid[k]`` contains the N-D array of indices along the
``k-th`` axis. Explicitly::
grid[k,i0,i1,...,iN-1] = ik
Examples
--------
>>> grid = np.indices((2, 3))
>>> grid.shape
(2, 2, 3)
>>> grid[0] # row indices
array([[0, 0, 0],
[1, 1, 1]])
>>> grid[1] # column indices
array([[0, 0, 0],
[1, 1, 1]], dtype=int32)
The indices can be used as an index into an array.
>>> x = np.arange(20).reshape(5, 4)
>>> row, col = np.indices((2, 3))
>>> x[row, col]
array([[0., 1., 2.],
[4., 5., 6.]])
Note that it would be more straightforward in the above example to
extract the required elements directly with ``x[:2, :3]``.
"""
return _mx_nd_np.indices(dimensions=dimensions, dtype=dtype, ctx=ctx)
72 changes: 71 additions & 1 deletion python/mxnet/symbol/numpy/_symbol.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@
from . import _internal as _npi

__all__ = ['zeros', 'ones', 'add', 'subtract', 'multiply', 'divide', 'mod', 'power', 'tensordot',
'linspace', 'expand_dims', 'tile', 'arange', 'split', 'concatenate', 'stack']
'linspace', 'expand_dims', 'tile', 'arange', 'split', 'concatenate', 'stack', 'indices']


def _num_outputs(sym):
Expand Down Expand Up @@ -1270,6 +1270,7 @@ def arange(start, stop=None, step=1, dtype=None, ctx=None):
return _npi.arange(start=start, stop=stop, step=step, dtype=dtype, ctx=ctx)


# pylint: disable=redefined-outer-name
@set_module('mxnet.symbol.numpy')
def split(ary, indices_or_sections, axis=0):
"""Split an array into multiple sub-arrays.
Expand Down Expand Up @@ -1310,6 +1311,7 @@ def split(ary, indices_or_sections, axis=0):
raise ValueError('indices_or_sections must either int or tuple of ints')
ret = _npi.split(ary, indices, axis, False, sections)
return ret
# pylint: enable=redefined-outer-name


@set_module('mxnet.symbol.numpy')
Expand Down Expand Up @@ -1362,4 +1364,72 @@ def get_list(arrays):
return _npi.stack(*arrays, axis=axis, out=out)


# pylint: disable=redefined-outer-name
@set_module('mxnet.symbol.numpy')
def indices(dimensions, dtype=_np.int32, ctx=None):
"""Return an array representing the indices of a grid.
Compute an array where the subarrays contain index values 0,1,...
varying only along the corresponding axis.
Parameters
----------
dimensions : sequence of ints
The shape of the grid.
dtype : data-type, optional
The desired data-type for the array. Default is `float32`.
ctx : device context, optional
Device context on which the memory is allocated. Default is
`mxnet.context.current_context()`.
Returns
-------
grid : _Symbol
The array of grid indices,
``grid.shape = (len(dimensions),) + tuple(dimensions)``.
Notes
-----
The output shape is obtained by prepending the number of dimensions
in front of the tuple of dimensions, i.e. if `dimensions` is a tuple
``(r0, ..., rN-1)`` of length ``N``, the output shape is
``(N,r0,...,rN-1)``.
The subarrays ``grid[k]`` contains the N-D array of indices along the
``k-th`` axis. Explicitly::
grid[k,i0,i1,...,iN-1] = ik
Examples
--------
>>> grid = np.indices((2, 3))
>>> grid.shape
(2, 2, 3)
>>> grid[0] # row indices
array([[0, 0, 0],
[1, 1, 1]])
>>> grid[1] # column indices
array([[0, 0, 0],
[1, 1, 1]], dtype=int32)
The indices can be used as an index into an array.
>>> x = np.arange(20).reshape(5, 4)
>>> row, col = np.indices((2, 3))
>>> x[row, col]
array([[0., 1., 2.],
[4., 5., 6.]])
Note that it would be more straightforward in the above example to
extract the required elements directly with ``x[:2, :3]``.
"""
if isinstance(dimensions, (tuple, list)):
if ctx is None:
ctx = current_context()
return _npi.indices(dimensions=dimensions, dtype=dtype, ctx=ctx)
else:
raise ValueError("The dimensions must be sequence of ints")
# pylint: enable=redefined-outer-name


_set_np_symbol_class(_Symbol)
33 changes: 33 additions & 0 deletions src/operator/numpy/np_init_op.cc
Original file line number Diff line number Diff line change
Expand Up @@ -24,10 +24,33 @@
*/
#include "../tensor/init_op.h"
#include "../tensor/elemwise_unary_op.h"
#include "./np_init_op.h"

namespace mxnet {
namespace op {

DMLC_REGISTER_PARAMETER(IndicesOpParam);

inline bool NumpyIndicesShape(const nnvm::NodeAttrs& attrs,
mxnet::ShapeVector* in_shapes,
mxnet::ShapeVector* out_shapes) {
const IndicesOpParam& param = nnvm::get<IndicesOpParam>(attrs.parsed);
CHECK_EQ(in_shapes->size(), 0U);
CHECK_EQ(out_shapes->size(), 1U);
CHECK_GE(param.dimensions.ndim(), 0)
<< "_npi_indices dimensions the number of dim must not be less than 0";
mxnet::TShape param_dim = param.dimensions;
if (!shape_is_known(param_dim)) return false;
const int indim = param.dimensions.ndim();
mxnet::TShape ret(indim + 1, -1);
ret[0] = indim;
for (int i = 1; i < indim + 1; ++i) {
ret[i] = param.dimensions[i-1];
}
SHAPE_ASSIGN_CHECK(*out_shapes, 0, ret);
return shape_is_known(out_shapes->at(0));
}

NNVM_REGISTER_OP(_npi_zeros)
.set_num_inputs(0)
.set_num_outputs(1)
Expand Down Expand Up @@ -110,5 +133,15 @@ NNVM_REGISTER_OP(_npi_arange)
.set_attr<FCompute>("FCompute<cpu>", RangeCompute<cpu, RangeParam>)
.add_arguments(RangeParam::__FIELDS__());

NNVM_REGISTER_OP(_npi_indices)
.describe("Return an array representing the indices of a grid.")
.set_num_inputs(0)
.set_num_outputs(1)
.set_attr_parser(ParamParser<IndicesOpParam>)
.set_attr<mxnet::FInferShape>("FInferShape", NumpyIndicesShape)
.set_attr<nnvm::FInferType>("FInferType", InitType<IndicesOpParam>)
.set_attr<FCompute>("FCompute<cpu>", IndicesCompute<cpu>)
.add_arguments(IndicesOpParam::__FIELDS__());

} // namespace op
} // namespace mxnet
4 changes: 4 additions & 0 deletions src/operator/numpy/np_init_op.cu
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
*/

#include "../tensor/init_op.h"
#include "./np_init_op.h"

namespace mxnet {
namespace op {
Expand All @@ -43,5 +44,8 @@ NNVM_REGISTER_OP(_np_ones_like)
NNVM_REGISTER_OP(_npi_arange)
.set_attr<FCompute>("FCompute<gpu>", RangeCompute<gpu, RangeParam>);

NNVM_REGISTER_OP(_npi_indices)
.set_attr<FCompute>("FCompute<gpu>", IndicesCompute<gpu>);

} // namespace op
} // namespace mxnet
Loading

0 comments on commit b37a698

Please sign in to comment.