This is an automated email from the ASF dual-hosted git repository.
jxie pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-mxnet.git
The following commit(s) were added to refs/heads/master by this push:
new 4efbbf5 [MXNET-176] Support axis=None for sum operator (#10405)
4efbbf5 is described below
commit 4efbbf5b281f26e9fd80db132012c83b135ce3a7
Author: Haibin Lin <linhaibin.e...@gmail.com>
AuthorDate: Fri Apr 6 10:41:59 2018 -0700
[MXNET-176] Support axis=None for sum operator (#10405)
* initial commit
* add doc
* update nnvm
* fix squeeze param
* fix gpu compliation
---
3rdparty/nnvm | 2 +-
cpp-package/scripts/OpWrapperGenerator.py | 1 +
src/operator/tensor/broadcast_reduce_op.h | 60 ++++++++++---------
src/operator/tensor/broadcast_reduce_op_value.cc | 4 +-
src/operator/tensor/broadcast_reduce_op_value.cu | 3 +-
src/operator/tensor/matrix_op.cc | 2 +-
src/operator/tensor/square_sum-inl.h | 30 ++++++----
tests/python/unittest/test_operator.py | 75 +++++++++++++-----------
8 files changed, 96 insertions(+), 81 deletions(-)
diff --git a/3rdparty/nnvm b/3rdparty/nnvm
index c342da7..2bc5144 160000
--- a/3rdparty/nnvm
+++ b/3rdparty/nnvm
@@ -1 +1 @@
-Subproject commit c342da72271c85e477480323f1d91997c6101ac0
+Subproject commit 2bc5144cd3733fd239287e3560c7db8285d21f02
diff --git a/cpp-package/scripts/OpWrapperGenerator.py
b/cpp-package/scripts/OpWrapperGenerator.py
index ac95773..0c000d9 100644
--- a/cpp-package/scripts/OpWrapperGenerator.py
+++ b/cpp-package/scripts/OpWrapperGenerator.py
@@ -94,6 +94,7 @@ class Arg:
'int or None':'dmlc::optional<int>',\
'long':'int64_t',\
'double':'double',\
+ 'Shape or None':'dmlc::optional<Shape>',\
'string':'const std::string&'}
name = ''
type = ''
diff --git a/src/operator/tensor/broadcast_reduce_op.h
b/src/operator/tensor/broadcast_reduce_op.h
index 2f6864b..1d4ef0a 100644
--- a/src/operator/tensor/broadcast_reduce_op.h
+++ b/src/operator/tensor/broadcast_reduce_op.h
@@ -37,11 +37,11 @@
namespace mxnet {
namespace op {
struct ReduceAxesParam : public dmlc::Parameter<ReduceAxesParam> {
- TShape axis;
+ dmlc::optional<TShape> axis;
bool keepdims;
bool exclude;
DMLC_DECLARE_PARAMETER(ReduceAxesParam) {
- DMLC_DECLARE_FIELD(axis).set_default(TShape())
+ DMLC_DECLARE_FIELD(axis).set_default(dmlc::optional<TShape>())
.describe(R"code(The axis or axes along which to perform the reduction.
The default, `axis=()`, will compute over all elements into a
@@ -66,12 +66,12 @@ struct ReduceAxesParam : public
dmlc::Parameter<ReduceAxesParam> {
struct NormParam : public dmlc::Parameter<NormParam> {
int ord;
- TShape axis;
+ dmlc::optional<TShape> axis;
bool keepdims;
DMLC_DECLARE_PARAMETER(NormParam) {
DMLC_DECLARE_FIELD(ord).set_default(2)
.describe("Order of the norm. Currently ord=2 is supported.");
- DMLC_DECLARE_FIELD(axis).set_default(TShape())
+ DMLC_DECLARE_FIELD(axis).set_default(dmlc::optional<TShape>())
.describe(R"code(The axis or axes along which to perform the reduction.
The default, `axis=()`, will compute over all elements into a
scalar array with shape `(1,)`.
@@ -198,17 +198,18 @@ inline bool ReduceAxisShape(const nnvm::NodeAttrs& attrs,
return true;
}
-inline TShape ReduceAxesShapeImpl(const TShape& ishape, const TShape& axis,
+inline TShape ReduceAxesShapeImpl(const TShape& ishape, const
dmlc::optional<TShape>& axis,
bool keepdims, bool exclude) {
- if (axis.ndim() == 0) {
+ // if axis doesn't have value, treat it same TShape().
+ if (!axis.has_value() || axis.value().ndim() == 0) {
if (keepdims) {
return TShape(ishape.ndim());
} else {
return TShape(1);
}
}
-
- TShape axes(axis);
+ // axis has value
+ TShape axes(axis.value());
for (index_t i = 0; i < axes.ndim(); i++) {
if (axes[i] < 0) {
axes[i] += ishape.ndim();
@@ -225,7 +226,7 @@ inline TShape ReduceAxesShapeImpl(const TShape& ishape,
const TShape& axis,
<< "Reduction axis " << axes[axes.ndim()-1]
<< " Exceeds input dimensions " << ishape;
CHECK_GE(axes[0], 0)
- << "Reduction axis " << axis
+ << "Reduction axis " << axis.value()
<< " Exceeds input dimensions " << ishape;
TShape oshape;
@@ -373,11 +374,12 @@ inline void BroadcastReduceShapeCompact(const TShape&
big, const TShape& small,
}
}
-inline bool SumOpForwardInferStorageType(const nnvm::NodeAttrs& attrs,
- const int dev_mask,
- DispatchMode* dispatch_mode,
- std::vector<int>* in_attrs,
- std::vector<int>* out_attrs) {
+// infer storage function for sum(csr) and mean(csr)
+inline bool ReduceAxesOpForwardStorage(const nnvm::NodeAttrs& attrs,
+ const int dev_mask,
+ DispatchMode* dispatch_mode,
+ std::vector<int>* in_attrs,
+ std::vector<int>* out_attrs) {
CHECK_EQ(in_attrs->size(), 1U);
CHECK_EQ(out_attrs->size(), 1U);
const ReduceAxesParam& param = nnvm::get<ReduceAxesParam>(attrs.parsed);
@@ -389,15 +391,14 @@ inline bool SumOpForwardInferStorageType(const
nnvm::NodeAttrs& attrs,
const auto dispatch_ex =
invalid_ctx ? DispatchMode::kFComputeFallback :
DispatchMode::kFComputeEx;
if (!dispatched && in_stype == kDefaultStorage) {
- // When input is dense output storage is set as dense and dispatched to
+ // When input is dense output storage is set as dense and dispatched to
// dense operator
dispatched = storage_type_assign(&out_stype, kDefaultStorage,
dispatch_mode,
DispatchMode::kFCompute);
}
-
- if (!dispatched && in_stype == kCSRStorage && param.axis.ndim() == 1 &&
- (param.axis[0] == 0 || param.axis[0] == 1) && !param.keepdims &&
- !param.exclude) {
+ TShape axis = param.axis.has_value() ? param.axis.value() : TShape();
+ if (!dispatched && in_stype == kCSRStorage && axis.ndim() == 1 &&
+ (axis[0] == 0 || axis[0] == 1) && !param.keepdims && !param.exclude) {
// If input is csr and axis is 0 or 1, and neither of keepdims or exclude
// are set, dipsatch to sparse operator and output storage is set as dense
dispatched = storage_type_assign(&out_stype, kDefaultStorage,
dispatch_mode,
@@ -695,12 +696,14 @@ template <typename xpu, typename red_op, bool normalize =
false>
void ReduceCsr(const nnvm::NodeAttrs& attrs, mshadow::Stream<xpu>* s, const
OpContext& ctx,
const NDArray& input, const OpReqType req, NDArray* output) {
const ReduceAxesParam& param = nnvm::get<ReduceAxesParam>(attrs.parsed);
- CHECK_EQ(param.axis.ndim(), 1U) << "sum(csr)/mean(csr) only supports axis 0
or 1";
- CHECK(param.axis[0] == 0 || param.axis[0] == 1)
- << "sum(csr)/mean(csr) only support axis 0 or 1";
+ CHECK(param.axis.has_value());
+ const TShape axis = param.axis.value();
+ CHECK_EQ(axis.ndim(), 1U) << "sum(csr)/mean(csr) only supports axis 0 or 1";
+ CHECK(axis[0] == 0 || axis[0] == 1)
+ << "sum(csr)/mean(csr) only support axis 0 or 1";
CHECK(!param.keepdims) << "keepdims not supported for sparse";
CHECK(!param.exclude) << "exclude not supported for sparse";
- ReduceCsrImpl<xpu, red_op, normalize>(s, ctx, input, req, output,
param.axis);
+ ReduceCsrImpl<xpu, red_op, normalize>(s, ctx, input, req, output, axis);
}
template <typename xpu, typename reducer, bool normalize = false>
@@ -1039,20 +1042,19 @@ void L2NormComputeEx(const nnvm::NodeAttrs& attrs,
CHECK_EQ(param.ord, 2) << "norm only support ord=2";
mshadow::Stream<xpu>* s = ctx.get_stream<xpu>();
const NDArrayStorageType istype = inputs[0].storage_type();
- if ((istype == kRowSparseStorage || istype == kCSRStorage)
- && param.axis.ndim() == 0) {
+ const TShape axis = param.axis.has_value() ? param.axis.value() : TShape();
+ if ((istype == kRowSparseStorage || istype == kCSRStorage) && axis.ndim() ==
0) {
// We only support norm on the entire array for now.
L2NormComputeSparseImpl<xpu>(s, inputs[0], req[0], outputs[0].data());
-
} else if (istype == kCSRStorage) {
CHECK_EQ(inputs[0].shape().ndim(), 2U)
<< "norm(csr) op only supports 2D ndarray as input";
- CHECK_EQ(param.axis.ndim(), 1U) << "sum(csr)/mean(csr) only supports axis
0 or 1";
- CHECK(param.axis[0] == 0 || param.axis[0] == 1)
+ CHECK_EQ(axis.ndim(), 1U) << "sum(csr)/mean(csr) only supports axis 0 or
1";
+ CHECK(axis[0] == 0 || axis[0] == 1)
<< "sum(csr)/mean(csr) only support axis 0 or 1";
CHECK(!param.keepdims) << "keepdims not supported for sparse";
NDArray output = outputs[0];
- ReduceCsrImpl<xpu, sq_sum, false>(s, ctx, inputs[0], req[0], &output,
param.axis);
+ ReduceCsrImpl<xpu, sq_sum, false>(s, ctx, inputs[0], req[0], &output,
axis);
CHECK_EQ(outputs[0].storage_type(), kDefaultStorage);
SqRootForL2<xpu>(ctx, req[0], outputs[0].data());
} else {
diff --git a/src/operator/tensor/broadcast_reduce_op_value.cc
b/src/operator/tensor/broadcast_reduce_op_value.cc
index da1d035..5f433cc 100644
--- a/src/operator/tensor/broadcast_reduce_op_value.cc
+++ b/src/operator/tensor/broadcast_reduce_op_value.cc
@@ -86,7 +86,7 @@ Example::
)code" ADD_FILELINE)
.set_attr<FCompute>("FCompute<cpu>", ReduceAxesCompute<cpu, mshadow::red::sum>)
.set_attr<FComputeEx>("FComputeEx<cpu>", SumOpForwardEx<cpu,
mshadow::red::sum>)
-.set_attr<FInferStorageType>("FInferStorageType", SumOpForwardInferStorageType)
+.set_attr<FInferStorageType>("FInferStorageType", ReduceAxesOpForwardStorage)
.set_attr<FResourceRequest>("FResourceRequest",
[](const NodeAttrs& attrs) {
return std::vector<ResourceRequest>{ResourceRequest::kTempSpace};
@@ -102,7 +102,7 @@ MXNET_ADD_SPARSE_OP_ALIAS(mean)
.describe(get_reduce_axes_description("mean", __LINE__))
.set_attr<FCompute>("FCompute<cpu>", ReduceAxesCompute<cpu, mshadow::red::sum,
true>)
.set_attr<FComputeEx>("FComputeEx<cpu>", SumOpForwardEx<cpu,
mshadow::red::sum, true>)
-.set_attr<FInferStorageType>("FInferStorageType", SumOpForwardInferStorageType)
+.set_attr<FInferStorageType>("FInferStorageType", ReduceAxesOpForwardStorage)
.set_attr<FResourceRequest>("FResourceRequest",
[](const NodeAttrs& attrs) {
return std::vector<ResourceRequest>{ResourceRequest::kTempSpace};
diff --git a/src/operator/tensor/broadcast_reduce_op_value.cu
b/src/operator/tensor/broadcast_reduce_op_value.cu
index dfff359..f1b19b4 100644
--- a/src/operator/tensor/broadcast_reduce_op_value.cu
+++ b/src/operator/tensor/broadcast_reduce_op_value.cu
@@ -39,9 +39,10 @@ void L2NormComputeEx<gpu>(const nnvm::NodeAttrs& attrs,
mshadow::Stream<gpu>* s = ctx.get_stream<gpu>();
const ReduceAxesParam& param = nnvm::get<ReduceAxesParam>(attrs.parsed);
const NDArrayStorageType in_stype = inputs[0].storage_type();
+ nnvm::TShape axis = param.axis.has_value() ? param.axis.value() : TShape();
// CSR and RowSparse only works on the entire array.
if ((in_stype == kCSRStorage || in_stype == kRowSparseStorage)
- && param.axis.ndim() == 0) {
+ && axis.ndim() == 0) {
L2NormComputeSparseImpl(s, inputs[0], req[0], outputs[0].data());
} else {
LogUnimplementedOp(attrs, ctx, inputs, req, outputs);
diff --git a/src/operator/tensor/matrix_op.cc b/src/operator/tensor/matrix_op.cc
index d751ca1..9037827 100644
--- a/src/operator/tensor/matrix_op.cc
+++ b/src/operator/tensor/matrix_op.cc
@@ -824,7 +824,7 @@ Examples::
.set_attr<FCompute>("FCompute<cpu>", UnaryOp::IdentityCompute<cpu>)
.set_attr<nnvm::FGradient>("FGradient",
ElemwiseGradUseNone{"_backward_squeeze"})
.add_argument("data", "NDArray-or-Symbol[]", "data to squeeze")
-.add_arguments(StackParam::__FIELDS__());
+.add_arguments(SqueezeParam::__FIELDS__());
NNVM_REGISTER_OP(_backward_squeeze)
.set_num_inputs(1)
diff --git a/src/operator/tensor/square_sum-inl.h
b/src/operator/tensor/square_sum-inl.h
index 0067ae1..162b3c8 100644
--- a/src/operator/tensor/square_sum-inl.h
+++ b/src/operator/tensor/square_sum-inl.h
@@ -53,13 +53,15 @@ inline bool SquareSumForwardInferStorageType(const
nnvm::NodeAttrs& attrs,
const auto& in_stype = in_attrs->at(0);
auto& out_stype = out_attrs->at(0);
bool dispatched = false;
- if (!dispatched && in_stype == kRowSparseStorage && param.axis[0] == 1 &&
param.keepdims) {
+ const TShape axis = param.axis.has_value() ? param.axis.value() : TShape();
+ if (!dispatched && in_stype == kRowSparseStorage &&
+ axis.ndim() > 0 && axis[0] == 1 && param.keepdims) {
// sum per row and keep dims
dispatched = storage_type_assign(&out_stype, kRowSparseStorage,
dispatch_mode, DispatchMode::kFComputeEx);
}
- if (!dispatched && in_stype == kRowSparseStorage &&
- (param.axis[0] == 0 || (param.axis[0] == 1 && !param.keepdims))) {
+ if (!dispatched && in_stype == kRowSparseStorage && axis.ndim() > 0 &&
+ (axis[0] == 0 || (axis[0] == 1 && !param.keepdims))) {
dispatched = storage_type_assign(&out_stype, kDefaultStorage,
dispatch_mode,
DispatchMode::kFComputeEx);
}
@@ -264,13 +266,15 @@ void SquareSumRspImpl(const nnvm::NodeAttrs& attrs,
NDArray* output) {
if (req == kNullOp) return;
const ReduceAxesParam& param = nnvm::get<ReduceAxesParam>(attrs.parsed);
- CHECK_EQ(param.axis.ndim(), 1U) << "_square_sum(row_sparse_matrix) only
supports axis=0 or 1";
- CHECK(param.axis[0] == 0 || param.axis[0] == 1)
+ CHECK(param.axis.has_value());
+ const TShape axis = param.axis.value();
+ CHECK_EQ(axis.ndim(), 1U) << "_square_sum(row_sparse_matrix) only supports
axis=0 or 1";
+ CHECK(axis[0] == 0 || axis[0] == 1)
<< "_square_sum(row_sparse_matrix) only supports axis=0 or 1";
CHECK_EQ(input.storage_type(), kRowSparseStorage)
<< "_square_sum op only supports row-sparse matrix as input";
int64_t out_data_size = 0;
- if (param.axis[0] == 0) { // axis = 0
+ if (axis[0] == 0) { // axis = 0
CHECK_EQ(output->storage_type(), kDefaultStorage);
out_data_size = input.storage_shape()[1];
} else if (param.keepdims) { // axis = 1, keepdims = true
@@ -305,7 +309,7 @@ void SquareSumRspImpl(const nnvm::NodeAttrs& attrs,
const int64_t nnr = input.storage_shape()[0];
const int64_t num_cols = input.storage_shape()[1];
const TBlob& in_data = input.data();
- if (0 == param.axis[0]) { // axis = 0, output is dense
+ if (0 == axis[0]) { // axis = 0, output is dense
MSHADOW_TYPE_SWITCH(out_data.type_flag_, DType, {
MXNET_ASSIGN_REQ_SWITCH(req, req_type, {
Kernel<SquareSumRspKernel<req_type, 0, false>, xpu>::Launch(s,
num_cols,
@@ -377,8 +381,10 @@ void SquareSumRspGradImpl(const nnvm::NodeAttrs& attrs,
NDArray* igrad) {
if (req == kNullOp) return;
const ReduceAxesParam& param = nnvm::get<ReduceAxesParam>(attrs.parsed);
- CHECK_EQ(param.axis.ndim(), 1U) << "_square_sum(row_sparse_matrix) only
supports axis=0/1";
- CHECK(param.axis[0] == 0 || param.axis[0] == 1)
+ CHECK(param.axis.has_value());
+ const TShape axis = param.axis.value();
+ CHECK_EQ(axis.ndim(), 1U) << "_square_sum(row_sparse_matrix) only supports
axis=0/1";
+ CHECK(axis[0] == 0 || axis[0] == 1)
<< "_square_sum(row_sparse_matrix) only supports axis=0 or 1";
CHECK(ograd.storage_type() == kDefaultStorage || ograd.storage_type() ==
kRowSparseStorage);
CHECK_EQ(input.storage_type(), kRowSparseStorage);
@@ -400,7 +406,7 @@ void SquareSumRspGradImpl(const nnvm::NodeAttrs& attrs,
igrad->CheckAndAlloc({input.aux_shape(rowsparse::kIdx)});
const TBlob& igrad_data = igrad->data();
const TBlob igrad_row_idx = igrad->aux_data(rowsparse::kIdx);
- if (0 == param.axis[0]) { // forward is sum per column
+ if (0 == axis[0]) { // forward is sum per column
MSHADOW_TYPE_SWITCH(igrad_data.type_flag_, DType, {
MSHADOW_IDX_TYPE_SWITCH(igrad_row_idx.type_flag_, IType, {
MXNET_ASSIGN_REQ_SWITCH(req, req_type, {
@@ -424,8 +430,8 @@ void SquareSumRspGradImpl(const nnvm::NodeAttrs& attrs,
})
}
} else if (ograd.storage_type() == kRowSparseStorage) {
- CHECK_EQ(1, param.axis[0]) << "SquareSumRspGradImpl only supports axis = 1"
- " when ograd_stype = kRowSparseStorage";
+ CHECK_EQ(1, axis[0]) << "SquareSumRspGradImpl only supports axis = 1"
+ " when ograd_stype = kRowSparseStorage";
CHECK_EQ(ograd.shape().ndim(), 2U);
const TBlob ograd_row_idx = ograd.aux_data(rowsparse::kIdx);
CHECK(ograd_row_idx.Size() == in_row_idx.Size() || in_row_idx.Size() ==
in_data.shape_[0]);
diff --git a/tests/python/unittest/test_operator.py
b/tests/python/unittest/test_operator.py
index 629304d..61b4478 100644
--- a/tests/python/unittest/test_operator.py
+++ b/tests/python/unittest/test_operator.py
@@ -1737,8 +1737,8 @@ def test_reshape():
@with_seed()
def test_reduce():
sample_num = 500
- def test_reduce_inner(numpy_reduce_func, numpy_reduce_grad_func,
mx_reduce_sym, nan_prob = 0,
- test_exclude = True):
+ def test_reduce_inner(numpy_reduce_func, numpy_reduce_grad_func,
mx_reduce_sym, nan_prob=0,
+ test_exclude=True, test_none_axis=False):
for i in range(sample_num):
# Generate random data that has ndim between 1-7 and all the shape
dims between 1-5
# Insert a NaN with probability equal to nan_prob
@@ -1763,7 +1763,10 @@ def test_reduce():
keepdims = np.random.randint(0, 2)
a = mx.symbol.Variable('a')
if axes is None:
- b = mx_reduce_sym(a, keepdims=keepdims)
+ if test_none_axis:
+ b = mx_reduce_sym(a, keepdims=keepdims, axis=axes)
+ else:
+ b = mx_reduce_sym(a, keepdims=keepdims)
elif exclude and isinstance(axes, tuple) and len(axes) < ndim:
naxes = [i for i in range(ndim) if i not in axes]
b = mx_reduce_sym(a, axis=naxes, keepdims=keepdims,
exclude=True)
@@ -1795,38 +1798,40 @@ def test_reduce():
equal_backward = almost_equal_ignore_nan(grad_nd.asnumpy(),
bc_grad_groundtruth, 1E-4, 1E-4)
assert equal_backward
- test_reduce_inner(lambda data, axis, keepdims:np_reduce(data, axis,
keepdims, np.sum),
- lambda outgrad, data, outdata, axis, keepdims,
keepdim_shape:
- outgrad.reshape(keepdim_shape),
- mx.symbol.sum)
- test_reduce_inner(lambda data, axis, keepdims:np_reduce(data, axis,
keepdims, np.mean),
- lambda outgrad, data, outdata, axis, keepdims,
keepdim_shape:
-
outgrad.reshape(keepdim_shape)/(data.size/outdata.size),
- mx.symbol.mean)
- test_reduce_inner(lambda data, axis, keepdims:np_reduce(data, axis,
keepdims, np.prod),
- lambda outgrad, data, outdata, axis, keepdims,
keepdim_shape:
- outgrad.reshape(keepdim_shape) *
(outdata.reshape(keepdim_shape) / data),
- mx.symbol.prod)
- test_reduce_inner(lambda data, axis, keepdims:np_reduce(data, axis,
keepdims, np.nansum),
- lambda outgrad, data, outdata, axis, keepdims,
keepdim_shape:
- np.where(np.isnan(data), 0,
outgrad.reshape(keepdim_shape)),
- mx.symbol.nansum, 0.3)
- test_reduce_inner(lambda data, axis, keepdims:np_reduce(data, axis,
keepdims, np.nanprod),
- lambda outgrad, data, outdata, axis, keepdims,
keepdim_shape:
- np.where(np.isnan(data), 0,
outgrad.reshape(keepdim_shape) * (outdata.reshape(keepdim_shape) / data)),
- mx.symbol.nanprod, 0.3)
- test_reduce_inner(lambda data, axis, keepdims:np_reduce(data, axis,
keepdims, np.max),
- lambda outgrad, data, outdata, axis, keepdims,
keepdim_shape:
- outgrad.reshape(keepdim_shape) * (np.equal(data,
outdata.reshape(keepdim_shape)).astype(np.float)),
- mx.symbol.max)
- test_reduce_inner(lambda data, axis, keepdims:np_reduce(data, axis,
keepdims, np.min),
- lambda outgrad, data, outdata, axis, keepdims,
keepdim_shape:
- outgrad.reshape(keepdim_shape) * (np.equal(data,
outdata.reshape(keepdim_shape)).astype(np.float)),
- mx.symbol.min)
- test_reduce_inner(lambda data, axis, keepdims:np_reduce(data, axis,
keepdims, np.linalg.norm),
- lambda outgrad, data, outdata, axis, keepdims,
keepdim_shape:
- outgrad.reshape(keepdim_shape) * (data /
outdata.reshape(keepdim_shape)),
- mx.symbol.norm, test_exclude=False)
+ test_none_axis = [True, False]
+ for test_none in test_none_axis:
+ test_reduce_inner(lambda data, axis, keepdims:np_reduce(data, axis,
keepdims, np.sum),
+ lambda outgrad, data, outdata, axis, keepdims,
keepdim_shape:
+ outgrad.reshape(keepdim_shape),
+ mx.symbol.sum, test_none_axis=test_none)
+ test_reduce_inner(lambda data, axis, keepdims:np_reduce(data, axis,
keepdims, np.mean),
+ lambda outgrad, data, outdata, axis, keepdims,
keepdim_shape:
+
outgrad.reshape(keepdim_shape)/(data.size/outdata.size),
+ mx.symbol.mean, test_none_axis=test_none)
+ test_reduce_inner(lambda data, axis, keepdims:np_reduce(data, axis,
keepdims, np.prod),
+ lambda outgrad, data, outdata, axis, keepdims,
keepdim_shape:
+ outgrad.reshape(keepdim_shape) *
(outdata.reshape(keepdim_shape) / data),
+ mx.symbol.prod, test_none_axis=test_none)
+ test_reduce_inner(lambda data, axis, keepdims:np_reduce(data, axis,
keepdims, np.nansum),
+ lambda outgrad, data, outdata, axis, keepdims,
keepdim_shape:
+ np.where(np.isnan(data), 0,
outgrad.reshape(keepdim_shape)),
+ mx.symbol.nansum, 0.3, test_none_axis=test_none)
+ test_reduce_inner(lambda data, axis, keepdims:np_reduce(data, axis,
keepdims, np.nanprod),
+ lambda outgrad, data, outdata, axis, keepdims,
keepdim_shape:
+ np.where(np.isnan(data), 0,
outgrad.reshape(keepdim_shape) * (outdata.reshape(keepdim_shape) / data)),
+ mx.symbol.nanprod, 0.3, test_none_axis=test_none)
+ test_reduce_inner(lambda data, axis, keepdims:np_reduce(data, axis,
keepdims, np.max),
+ lambda outgrad, data, outdata, axis, keepdims,
keepdim_shape:
+ outgrad.reshape(keepdim_shape) * (np.equal(data,
outdata.reshape(keepdim_shape)).astype(np.float)),
+ mx.symbol.max, test_none_axis=test_none)
+ test_reduce_inner(lambda data, axis, keepdims:np_reduce(data, axis,
keepdims, np.min),
+ lambda outgrad, data, outdata, axis, keepdims,
keepdim_shape:
+ outgrad.reshape(keepdim_shape) * (np.equal(data,
outdata.reshape(keepdim_shape)).astype(np.float)),
+ mx.symbol.min, test_none_axis=test_none)
+ test_reduce_inner(lambda data, axis, keepdims:np_reduce(data, axis,
keepdims, np.linalg.norm),
+ lambda outgrad, data, outdata, axis, keepdims,
keepdim_shape:
+ outgrad.reshape(keepdim_shape) * (data /
outdata.reshape(keepdim_shape)),
+ mx.symbol.norm, test_exclude=False,
test_none_axis=test_none)
@with_seed()
--
To stop receiving notification emails like this one, please contact
j...@apache.org.
