Numpy indexing

python/mxnet/ndarray/numpy/_op.py

@@ -28,12 +28,12 @@
 from . import _internal as _npi
 from ..ndarray import NDArray
 
-__all__ = ['zeros', 'ones', 'maximum', 'minimum', 'stack', 'arange', 'argmax',
+__all__ = ['zeros', 'ones', 'full', 'maximum', 'minimum', 'stack', 'arange', 'argmax',
            'add', 'subtract', 'multiply', 'divide', 'mod', 'power', 'concatenate',
            'clip', 'split', 'swapaxes', 'expand_dims', 'tile', 'linspace', 'eye',
            'sin', 'cos', 'sinh', 'cosh', 'log10', 'sqrt', 'abs', 'exp', 'arctan', 'sign', 'log',
            'degrees', 'log2', 'rint', 'radians', 'mean', 'reciprocal', 'square', 'arcsin',
-           'argsort', 'hstack', 'tensordot']
+           'argsort', 'hstack', 'tensordot', "slice_assign_scalar", "slice_assign"]
 
 
 @set_module('mxnet.ndarray.numpy')
@@ -178,6 +178,62 @@ def ones(shape, dtype=None, **kwargs):
     return _npi.ones(shape=shape, ctx=ctx, dtype=dtype, **kwargs)
 
 
+@set_module('mxnet.ndarray.numpy')
+def full(shape, fill_value, dtype=None, ctx=None, out=None):
+    """
+    Return a new array of given shape and type, filled with `fill_value`.
+
+    Parameters
+    ----------
+    shape : int or sequence of ints
+        Shape of the new array, e.g., ``(2, 3)`` or ``2``.
+    fill_value : scalar
+        Fill value.
+    dtype : data-type, optional
+        The desired data-type for the array. The default, `None`, means
+        `np.array(fill_value).dtype`.
+    ctx: to specify the device, e.g. the i-th GPU.
+    out : ndarray or None, optional
+        A location into which the result is stored.
+        If provided, it must have the same shape and dtype as input ndarray.
+        If not provided or `None`, a freshly-allocated array is returned.
+
+    Returns
+    -------
+    out : ndarray
+        Array of `fill_value` with the given shape, dtype, and order.
+
+    Notes
+    -----
+    This function differs from the original `numpy.full
+    https://docs.scipy.org/doc/numpy/reference/generated/numpy.full.html`_ in
+    the following way(s):
+
+    - Have an additional `ctx` argument to specify the device
+    - Have an additional `out` argument
+    - Currently does not support `order` selection
+
+    See Also
+    --------
+    empty : Return a new uninitialized array.
+    ones : Return a new array setting values to one.
+    zeros : Return a new array setting values to zero.
+
+    Examples
+    --------
+    >>> np.full((2, 2), 10)
+    array([[10., 10.],
+           [10., 10.]])
+    >>> np.full((2, 2), 2, dtype=np.int32, ctx=mx.cpu(0))
+    array([[2, 2],
+           [2, 2]], dtype=int32)
+    """
+    if ctx is None:
+        ctx = current_context()
+    dtype = _np.float32 if dtype is None else dtype
+    return _npi.full(shape=shape, value=fill_value, ctx=ctx, dtype=dtype, out=out)
+
+
 #pylint: disable= too-many-arguments, no-member, protected-access
 def _ufunc_helper(lhs, rhs, fn_array, fn_scalar, lfn_scalar, rfn_scalar=None, out=None):
     """ Helper function for element-wise operation.
@@ -1902,3 +1958,21 @@ def arcsin(x, out=None, **kwargs):
     http://www.math.sfu.ca/~cbm/aands/
     """
     return _unary_func_helper(x, _npi.arcsin, _np.arcsin, out=out, **kwargs)
+
+@set_module('mxnet.ndarray.numpy')
+def slice_assign_scalar(lhs, value, begin, end, step, out=None):  # pylint: disable=too-many-arguments
+    """
+    Assign the rhs to a cropped subset of this ndarray in place. Returns the view of this ndarray.
+    `rhs` and `lhs` should be of the same data type, and on the same device.
+    The shape of rhs should be the same as the cropped shape of this ndarray.
+    """
+    _npi.slice_assign_scalar(lhs, value, begin=begin, end=end, step=step, out=out)
+
+
+@set_module('mxnet.ndarray.numpy')
+def slice_assign(lhs, rhs, begin, end, step, out=None):  # pylint: disable=too-many-arguments
+    """
+    Assign the scalar to a cropped subset of this ndarray. Value will broadcast to the shape of the
+    cropped shape and will be cast to the same dtype of the ndarray.
+    """
+    return _npi.slice_assign(lhs, rhs, begin=begin, end=end, step=step, out=out)