add np.product (apache#17489)

anirudh2290 · May 29, 2020 · 146fa60 · 146fa60
1 parent 58dcd41
commit 146fa60
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 1 deletion.
diff --git a/python/mxnet/_numpy_op_doc.py b/python/mxnet/_numpy_op_doc.py
@@ -63,7 +63,7 @@ def _np_all(a, axis=None, keepdims=False, out=None):
  >>> o=np.array(False)
  >>> z=np.all([-1, 4, 5], out=o)
  >>> id(z), id(o), z
- (28293632, 28293632, array(True)) # may vary 
+ (28293632, 28293632, array(True)) # may vary
  """
  pass
 
@@ -1171,6 +1171,17 @@ def _np_prod(a, axis=None, dtype=None, out=None, keepdims=False):
  pass
 
 
+def _np_product(a, axis=None, dtype=None, out=None, keepdims=False):
+ """
+ Return the product of array elements over a given axis.
+
+ See Also
+ --------
+ prod : equivalent function; see for details.
+ """
+ pass
+
+
 def _np_moveaxis(a, source, destination):
  """Move axes of an array to new positions.
  Other axes remain in their original order.

diff --git a/python/mxnet/numpy_dispatch_protocol.py b/python/mxnet/numpy_dispatch_protocol.py
@@ -119,6 +119,7 @@ def _run_with_array_ufunc_proto(*args, **kwargs):
  'atleast_2d',
  'atleast_3d',
  'prod',
+ 'product',
  'ravel',
  'repeat',
  'reshape',

diff --git a/src/operator/numpy/np_broadcast_reduce_op_value.cc b/src/operator/numpy/np_broadcast_reduce_op_value.cc
@@ -224,6 +224,7 @@ NNVM_REGISTER_OP(_backward_np_min)
 .set_attr<FCompute>("FCompute<cpu>", NumpyReduceAxesNoDTypeBackward<cpu, mshadow_op::eq>);
 
 NNVM_REGISTER_OP(_np_prod)
+.add_alias("_np_product")
 .set_num_inputs(1)
 .set_num_outputs(1)
 .set_attr_parser(ParamParser<NumpyReduceAxesParam>)

diff --git a/tests/python/unittest/test_numpy_interoperability.py b/tests/python/unittest/test_numpy_interoperability.py
@@ -1021,6 +1021,10 @@ def _add_workload_prod(array_pool):
  OpArgMngr.add_workload('prod', array_pool['4x1'])
 
 
+def _add_workload_product(array_pool):
+ OpArgMngr.add_workload('product', array_pool['4x1'])
+
+
 def _add_workload_repeat(array_pool):
  OpArgMngr.add_workload('repeat', array_pool['4x1'], 3)
  OpArgMngr.add_workload('repeat', np.array(_np.arange(12).reshape(4, 3)[:, 2]), 3)
@@ -2723,6 +2727,7 @@ def _prepare_workloads():
  _add_workload_ones_like(array_pool)
  _add_workload_atleast_nd()
  _add_workload_prod(array_pool)
+ _add_workload_product(array_pool)
  _add_workload_repeat(array_pool)
  _add_workload_reshape()
  _add_workload_rint(array_pool)