add np.product

apache · Feb 15, 2020 · e727d92 · e727d92
1 parent b6b1de0
commit e727d92
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
@@ -219,6 +219,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
@@ -998,6 +998,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)
@@ -1992,6 +1996,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)