apeforest commented on a change in pull request #15169: Softmax with length
URL: https://github.com/apache/incubator-mxnet/pull/15169#discussion_r304627396
##########
File path: src/operator/nn/softmax-inl.h
##########
@@ -464,35 +755,56 @@ void SoftmaxGradCompute(const nnvm::NodeAttrs& attrs,
mxnet::TShape shape = AxisShapeCompact(inputs[0].shape_, &axis, true);
int out_idx = softmax_has_dtype_override(attrs) ? 2 : 1;
+ out_idx = softmax_use_length(attrs) ? 3 : out_idx;
bool safe_acc = dmlc::GetEnv("MXNET_SAFE_ACCUMULATION", false);
MXNET_REAL_ACC_TYPE_SWITCH(inputs[0].type_flag_, OType, AType, {
MSHADOW_REAL_TYPE_SWITCH(outputs[0].type_flag_, DType, {
MXNET_ASSIGN_REQ_SWITCH(req[0], Req, {
- if (safe_acc) {
- if (shape.ndim() == 2) {
- SoftmaxGrad<OP1, OP2, Req, negate, AType>(
- ctx.get_stream<xpu>(), inputs[out_idx].dptr<OType>(),
- inputs[0].dptr<OType>(), outputs[0].dptr<DType>(),
- shape.get<2>(), axis, static_cast<DType>(temperature));
+ if (!softmax_use_length(attrs)) {
+ if (safe_acc) {
+ if (shape.ndim() == 2) {
+ SoftmaxGrad<OP1, OP2, Req, negate, AType>(
+ ctx.get_stream<xpu>(), inputs[out_idx].dptr<OType>(),
+ inputs[0].dptr<OType>(), outputs[0].dptr<DType>(),
+ shape.get<2>(), axis, static_cast<DType>(temperature));
+ } else {
+ SoftmaxGrad<OP1, OP2, Req, negate, AType>(
+ ctx.get_stream<xpu>(), inputs[out_idx].dptr<OType>(),
+ inputs[0].dptr<OType>(), outputs[0].dptr<DType>(),
+ shape.get<3>(), axis, static_cast<DType>(temperature));
+ }
} else {
- SoftmaxGrad<OP1, OP2, Req, negate, AType>(
- ctx.get_stream<xpu>(), inputs[out_idx].dptr<OType>(),
- inputs[0].dptr<OType>(), outputs[0].dptr<DType>(),
- shape.get<3>(), axis, static_cast<DType>(temperature));
+ if (shape.ndim() == 2) {
+ SoftmaxGrad<OP1, OP2, Req, negate, DType>(
+ ctx.get_stream<xpu>(), inputs[out_idx].dptr<OType>(),
+ inputs[0].dptr<OType>(), outputs[0].dptr<DType>(),
+ shape.get<2>(), axis, static_cast<DType>(temperature));
+ } else {
+ SoftmaxGrad<OP1, OP2, Req, negate, DType>(
+ ctx.get_stream<xpu>(), inputs[out_idx].dptr<OType>(),
+ inputs[0].dptr<OType>(), outputs[0].dptr<DType>(),
+ shape.get<3>(), axis, static_cast<DType>(temperature));
+ }
}
} else {
- if (shape.ndim() == 2) {
- SoftmaxGrad<OP1, OP2, Req, negate, DType>(
+ MXNET_INT_TYPE_SWITCH(inputs[2].type_flag_, IType, {
Review comment:
Are you referring to these two kernel functions?
https://github.com/apache/incubator-mxnet/blob/50c1c7ca7f2e6a7864e7a0aedc775ae7dd8be091/src/operator/nn/softmax-inl.h#L283
https://github.com/apache/incubator-mxnet/blob/50c1c7ca7f2e6a7864e7a0aedc775ae7dd8be091/src/operator/nn/softmax-inl.h#L339
They also only differ by a few lines. Would you think a better way to
consolidate them?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
