[FFI] part5: npx.batch_dot, npx.arange_like, npx.broadcast_like

python/mxnet/base.py

@@ -798,7 +798,8 @@ def write_all_str(module_file, module_all_list):
                               '_npx_masked_log_softmax', '_npx_activation',
                               '_npx_batch_norm', '_npx_fully_connected', '_npx_pick',
                               '_npx_convolution', '_npx_deconvolution', '_npx_pooling',
-                              '_npx_dropout', '_npx_one_hot', '_npx_rnn'}
+                              '_npx_dropout', '_npx_one_hot', '_npx_rnn', '_npx_batch_dot',
+                              '_npx_broadcast_like', '_npx_arange_like'}
 
 _NP_INTERNAL_OP_PREFIX = '_npi_'
 

python/mxnet/ndarray/numpy_extension/_op.py

@@ -26,7 +26,8 @@
 
 __all__ = ['softmax', 'log_softmax', 'masked_softmax', 'masked_log_softmax',
            'activation', 'batch_norm', 'fully_connected', 'pick', 'convolution',
-           'deconvolution', 'pooling', 'dropout', 'one_hot', 'rnn']
+           'deconvolution', 'pooling', 'dropout', 'one_hot', 'rnn',
+           'batch_dot', 'broadcast_like', 'arange_like']
 
 
 # pylint: disable=too-many-arguments
@@ -1033,3 +1034,132 @@ def rnn(data=None, parameters=None, state=None, state_cell=None, sequence_length
                                      state_size, num_layers, bidirectional, state_outputs,
                                      mode, p, use_sequence_length, projection_size,
                                      lstm_state_clip_min, lstm_state_clip_max, lstm_state_clip_nan)
+
+
+# pylint: disable=too-many-arguments
+@set_module('mxnet.ndarray.numpy_extension')
+def batch_dot(a, b, transpose_a=False, transpose_b=False, forward_stype="default"):
+    r"""Batchwise dot product.
+
+    ``batch_dot`` is used to compute dot product of ``x`` and ``y`` when ``x`` and
+    ``y`` are data in batch, namely N-D (N >= 3) arrays in shape of `(B0, ..., B_i, :, :)`.
+
+    For example, given ``x`` with shape `(B_0, ..., B_i, N, M)` and ``y`` with shape
+    `(B_0, ..., B_i, M, K)`, the result array will have shape `(B_0, ..., B_i, N, K)`,
+    which is computed by::
+
+       batch_dot(x,y)[b_0, ..., b_i, :, :] = dot(x[b_0, ..., b_i, :, :], y[b_0, ..., b_i, :, :])
+
+    Parameters
+    ----------
+    lhs : NDArray
+        The first input
+    rhs : NDArray
+        The second input
+    transpose_a : boolean, optional, default=0
+        If true then transpose the first input before dot.
+    transpose_b : boolean, optional, default=0
+        If true then transpose the second input before dot.
+    forward_stype : {None, 'csr', 'default', 'row_sparse'},optional, default='None'
+        The desired storage type of the forward output given by user,
+        if thecombination of input storage types and this hint does not matchany implemented ones,
+        the dot operator will perform fallback operationand still produce
+        an output of the desired storage type.
+
+    Returns
+    -------
+    out : NDArray or list of NDArrays
+        The output of this function.
+    """
+    return _api_internal.batch_dot(a, b, transpose_a, transpose_b, forward_stype)
+
+
+# pylint: disable=too-many-arguments
+@set_module('mxnet.ndarray.numpy_extension')
+def broadcast_like(lhs, rhs, lhs_axes=None, rhs_axes=None):
+    r"""Broadcasts lhs to have the same shape as rhs.
+
+    Broadcasting is a mechanism that allows NDArrays to perform arithmetic operations
+    with arrays of different shapes efficiently without creating multiple copies of arrays.
+    Also see, `Broadcasting <https://docs.scipy.org/doc/numpy/user/basics.broadcasting.html>`_
+    for more explanation.
+
+    Broadcasting is allowed on axes with size 1, such as from `(2,1,3,1)` to
+    `(2,8,3,9)`. Elements will be duplicated on the broadcasted axes.
+
+    Parameters
+    ----------
+    lhs : NDArray
+        First input.
+    rhs : NDArray
+        Second input.
+    lhs_axes : Shape or None, optional, default=None
+        Axes to perform broadcast on in the first input array
+    rhs_axes : Shape or None, optional, default=None
+        Axes to copy from the second input array
+
+    Returns
+    -------
+    out : NDArray or list of NDArrays
+        The output of this function.
+
+    example
+    -------
+    >>> a = np.array([[1,2,3]])
+    >>> b = np.array([[5,6,7],[7,8,9]])
+    >>> npx.broadcast_like(a, b)
+    array([[1., 2., 3.],
+           [1., 2., 3.]])
+    >>> a = np.array([9])
+    >>> b = np.array([1,2,3,4,5])
+    >>> npx.broadcast_like(a, b, lhs_axes=(0,), rhs_axes=(-1,))
+    array([9., 9., 9., 9., 9.])
+    """
+    return _api_internal.broadcast_like(lhs, rhs, lhs_axes, rhs_axes)
+
+
+# pylint: disable=too-many-arguments
+@set_module('mxnet.ndarray.numpy_extension')
+def arange_like(data, start=0.0, step=1.0, repeat=1, ctx=None, axis=None):
+    r"""Return an array with evenly spaced values. If axis is not given, the output will
+    have the same shape as the input array. Otherwise, the output will be a 1-D array with size of
+    the specified axis in input shape.
+
+    Parameters
+    ----------
+    data : NDArray
+        The input
+    start : double, optional, default=0
+        Start of interval. The interval includes this value. The default start value is 0.
+    step : double, optional, default=1
+        Spacing between values.
+    repeat : int, optional, default='1'
+        The repeating time of all elements.
+        E.g repeat=3, the element a will be repeated three times --> a, a, a.
+    ctx : string, optional, default=''
+        Context of output, in format [cpu|gpu|cpu_pinned](n).Only used for imperative calls.
+    axis : int or None, optional, default='None'
+        Arange elements according to the size of a certain axis of input array.
+        The negative numbers are interpreted counting from the backward.
+        If not provided, will arange elements according to the input shape.
+
+    Returns
+    -------
+    out : NDArray or list of NDArrays
+        The output of this function.
+
+    Example
+    -------
+    >>> x = np.random.uniform(0, 1, size=(3,4))
+    >>> x
+    array([[0.5488135 , 0.5928446 , 0.71518934, 0.84426576],
+           [0.60276335, 0.8579456 , 0.5448832 , 0.8472517 ],
+           [0.4236548 , 0.6235637 , 0.6458941 , 0.3843817 ]])
+    >>> npx.arange_like(x, start=0)
+    array([[ 0.,  1.,  2.,  3.],
+           [ 4.,  5.,  6.,  7.],
+           [ 8.,  9., 10., 11.]])
+    >>> npx.arange_like(x, start=0, axis=-1)
+    array([0., 1., 2., 3.])
+    """
+    return _api_internal.arange_like(data, start, step, repeat, ctx, axis)

python/mxnet/numpy_extension/_op.py

@@ -23,7 +23,8 @@
 
 __all__ = ['softmax', 'log_softmax', 'masked_softmax', 'masked_log_softmax',
            'activation', 'batch_norm', 'fully_connected', 'pick', 'convolution',
-           'deconvolution', 'pooling', 'dropout', 'one_hot', 'rnn']
+           'deconvolution', 'pooling', 'dropout', 'one_hot', 'rnn',
+           'batch_dot', 'broadcast_like', 'arange_like']
 
 
 # pylint: disable=too-many-arguments
@@ -965,3 +966,134 @@ def rnn(data=None, parameters=None, state=None, state_cell=None, sequence_length
                           projection_size=projection_size, lstm_state_clip_min=lstm_state_clip_min,
                           lstm_state_clip_max=lstm_state_clip_max,
                           lstm_state_clip_nan=lstm_state_clip_nan)
+
+
+# pylint: disable=too-many-arguments, unused-argument
+@set_module('mxnet.numpy_extension')
+def batch_dot(a, b, transpose_a=False, transpose_b=False, forward_stype="default"):
+    r"""Batchwise dot product.
+
+    ``batch_dot`` is used to compute dot product of ``x`` and ``y`` when ``x`` and
+    ``y`` are data in batch, namely N-D (N >= 3) arrays in shape of `(B0, ..., B_i, :, :)`.
+
+    For example, given ``x`` with shape `(B_0, ..., B_i, N, M)` and ``y`` with shape
+    `(B_0, ..., B_i, M, K)`, the result array will have shape `(B_0, ..., B_i, N, K)`,
+    which is computed by::
+
+       batch_dot(x,y)[b_0, ..., b_i, :, :] = dot(x[b_0, ..., b_i, :, :], y[b_0, ..., b_i, :, :])
+
+    Parameters
+    ----------
+    lhs : NDArray
+        The first input
+    rhs : NDArray
+        The second input
+    transpose_a : boolean, optional, default=0
+        If true then transpose the first input before dot.
+    transpose_b : boolean, optional, default=0
+        If true then transpose the second input before dot.
+    forward_stype : {None, 'csr', 'default', 'row_sparse'},optional, default='None'
+        The desired storage type of the forward output given by user,
+        if thecombination of input storage types and this hint does not matchany implemented ones,
+        the dot operator will perform fallback operationand still produce
+        an output of the desired storage type.
+
+    Returns
+    -------
+    out : NDArray or list of NDArrays
+        The output of this function.
+    """
+    return _mx_nd_npx.batch_dot(a=a, b=b, transpose_a=transpose_a,
+                                transpose_b=transpose_b, forward_stype=forward_stype)
+
+
+# pylint: disable=too-many-arguments, unused-argument
+@set_module('mxnet.numpy_extension')
+def broadcast_like(lhs, rhs, lhs_axes=None, rhs_axes=None):
+    r"""Broadcasts lhs to have the same shape as rhs.
+
+    Broadcasting is a mechanism that allows NDArrays to perform arithmetic operations
+    with arrays of different shapes efficiently without creating multiple copies of arrays.
+    Also see, `Broadcasting <https://docs.scipy.org/doc/numpy/user/basics.broadcasting.html>`_
+    for more explanation.
+
+    Broadcasting is allowed on axes with size 1, such as from `(2,1,3,1)` to
+    `(2,8,3,9)`. Elements will be duplicated on the broadcasted axes.
+
+    Parameters
+    ----------
+    lhs : NDArray
+        First input.
+    rhs : NDArray
+        Second input.
+    lhs_axes : Shape or None, optional, default=None
+        Axes to perform broadcast on in the first input array
+    rhs_axes : Shape or None, optional, default=None
+        Axes to copy from the second input array
+
+    Returns
+    -------
+    out : NDArray or list of NDArrays
+        The output of this function.
+
+    example
+    -------
+    >>> a = np.array([[1,2,3]])
+    >>> b = np.array([[5,6,7],[7,8,9]])
+    >>> npx.broadcast_like(a, b)
+    array([[1., 2., 3.],
+           [1., 2., 3.]])
+    >>> a = np.array([9])
+    >>> b = np.array([1,2,3,4,5])
+    >>> npx.broadcast_like(a, b, lhs_axes=(0,), rhs_axes=(-1,))
+    array([9., 9., 9., 9., 9.])
+    """
+    return _mx_nd_npx.broadcast_like(lhs=lhs, rhs=rhs, lhs_axes=lhs_axes, rhs_axes=rhs_axes)
+
+
+# pylint: disable=too-many-arguments, unused-argument
+@set_module('mxnet.numpy_extension')
+def arange_like(data, start=0.0, step=1.0, repeat=1, ctx=None, axis=None):
+    r"""Return an array with evenly spaced values. If axis is not given, the output will
+    have the same shape as the input array. Otherwise, the output will be a 1-D array with size of
+    the specified axis in input shape.
+
+    Parameters
+    ----------
+    data : NDArray
+        The input
+    start : double, optional, default=0
+        Start of interval. The interval includes this value. The default start value is 0.
+    step : double, optional, default=1
+        Spacing between values.
+    repeat : int, optional, default='1'
+        The repeating time of all elements.
+        E.g repeat=3, the element a will be repeated three times --> a, a, a.
+    ctx : string, optional, default=''
+        Context of output, in format [cpu|gpu|cpu_pinned](n).Only used for imperative calls.
+    axis : int or None, optional, default='None'
+        Arange elements according to the size of a certain axis of input array.
+        The negative numbers are interpreted counting from the backward.
+        If not provided, will arange elements according to the input shape.
+
+    Returns
+    -------
+    out : NDArray or list of NDArrays
+        The output of this function.
+
+    Example
+    -------
+    >>> x = np.random.uniform(0, 1, size=(3,4))
+    >>> x
+    array([[0.5488135 , 0.5928446 , 0.71518934, 0.84426576],
+           [0.60276335, 0.8579456 , 0.5448832 , 0.8472517 ],
+           [0.4236548 , 0.6235637 , 0.6458941 , 0.3843817 ]])
+    >>> npx.arange_like(x, start=0)
+    array([[ 0.,  1.,  2.,  3.],
+           [ 4.,  5.,  6.,  7.],
+           [ 8.,  9., 10., 11.]])
+    >>> npx.arange_like(x, start=0, axis=-1)
+    array([0., 1., 2., 3.])
+    """
+    return _mx_nd_npx.arange_like(data=data, start=start, step=step, repeat=repeat,
+                                  ctx=ctx, axis=axis)

src/api/operator/numpy_extension/npx_arange_like_op.cc

@@ -0,0 +1,78 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ * 
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ * 
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file npx_arange_like_op.cc
+ * \brief Implementation of the API of functions in src/operator/numpy_extension/npx_arange_like_op.cc
+ */
+#include <mxnet/api_registry.h>
+#include <mxnet/runtime/packed_func.h>
+#include "../utils.h"
+#include "../../../operator/tensor/init_op.h"
+
+namespace mxnet {
+
+MXNET_REGISTER_API("_npx.arange_like")
+.set_body([](runtime::MXNetArgs args, runtime::MXNetRetValue* ret) {
+  using namespace runtime;
+  nnvm::NodeAttrs attrs;
+  const nnvm::Op* op = Op::Get("_npx_arange_like");
+  op::RangeLikeParam param;
+  // inputs
+  int num_inputs = 1;
+  NDArray* inputs[] = {args[0].operator mxnet::NDArray*()};
+  // start
+  if (args[1].type_code() == kNull) {
+    param.start = 0.0;
+  } else {
+    param.start = args[1].operator double();
+  }
+  // step
+  if (args[2].type_code() == kNull) {
+    param.step = 1.0;
+  } else {
+    param.step = args[2].operator double();
+  }
+  // repeat
+  if (args[3].type_code() == kNull) {
+    param.repeat = 1;
+  } else {
+    param.repeat = args[3].operator int();
+  }
+  // ctx
+  if (args[4].type_code() != kNull) {
+    attrs.dict["ctx"] = args[4].operator std::string();
+  }
+  // axis
+  if (args[5].type_code() == kNull) {
+    param.axis = dmlc::nullopt;
+  } else {
+    param.axis = args[5].operator int();
+  }
+  attrs.op = op;
+  attrs.parsed = param;
+  SetAttrDict<op::RangeLikeParam>(&attrs);
+
+  // outputs
+  int num_outputs = 0;
+  auto ndoutputs = Invoke(op, &attrs, num_inputs, inputs, &num_outputs, nullptr);
+  *ret = reinterpret_cast<mxnet::NDArray*>(ndoutputs[0]);
+});
+
+}  // namespace mxnet