[GitHub] szha closed pull request #11335: Fix axis Bug in MKLDNN Softmax

[GitBox <gi...@apache.org>]

szha closed pull request #11335: Fix axis Bug in MKLDNN Softmax
URL: https://github.com/apache/incubator-mxnet/pull/11335
 
 
   

This is a PR merged from a forked repository.
As GitHub hides the original diff on merge, it is displayed below for
the sake of provenance:

As this is a foreign pull request (from a fork), the diff is supplied
below (as it won't show otherwise due to GitHub magic):

diff --git a/src/operator/nn/mkldnn/mkldnn_softmax.cc b/src/operator/nn/mkldnn/mkldnn_softmax.cc
index aa59f13d06d..acfa358a796 100644
--- a/src/operator/nn/mkldnn/mkldnn_softmax.cc
+++ b/src/operator/nn/mkldnn/mkldnn_softmax.cc
@@ -26,6 +26,7 @@
 #include "../softmax-inl.h"
 #include "./mkldnn_ops-inl.h"
 #include "./mkldnn_base-inl.h"
+#include "../../tensor/broadcast_reduce_op.h"
 
 #if MXNET_USE_MKLDNN == 1
 namespace mxnet {
@@ -38,11 +39,13 @@ void MKLDNNSoftmaxForward(const nnvm::NodeAttrs& attrs, const OpContext &ctx,
   auto input_mem = in_data.GetMKLDNNData();
   mkldnn::memory::primitive_desc data_mpd = input_mem->get_primitive_desc();
   mkldnn::memory::desc data_md = data_mpd.desc();
+  int axis = CheckAxis(param.axis, in_data.shape().ndim());
+
   auto cpu_engine = data_mpd.get_engine();
   auto prop = ctx.is_train
     ? mkldnn::prop_kind::forward_training : mkldnn::prop_kind::forward_scoring;
   mkldnn::softmax_forward::desc desc = mkldnn::softmax_forward::desc(prop,
-      data_md, param.axis);
+      data_md, axis);
   mkldnn::softmax_forward::primitive_desc pdesc(desc, cpu_engine);
 
   auto output_memory = out_data.GetMKLDNNData();
diff --git a/src/operator/nn/softmax.cc b/src/operator/nn/softmax.cc
index f8cc6fee9a2..e9b104f1286 100644
--- a/src/operator/nn/softmax.cc
+++ b/src/operator/nn/softmax.cc
@@ -38,10 +38,8 @@ static void SoftmaxComputeExCPU(const nnvm::NodeAttrs& attrs,
                                 const std::vector<NDArray>& inputs,
                                 const std::vector<OpReqType>& req,
                                 const std::vector<NDArray>& outputs) {
-  const SoftmaxParam& param = nnvm::get<SoftmaxParam>(attrs.parsed);
   // It seems MKLDNN softmax doesn't support training.
-  // and it only supports non-negative axis.
-  if (SupportMKLDNN(inputs[0]) && !ctx.is_train && param.axis >= 0) {
+  if (SupportMKLDNN(inputs[0]) && !ctx.is_train) {
     MKLDNN_OPCHECK_INIT(false, outputs.size(), inputs, outputs);
     MKLDNNSoftmaxForward(attrs, ctx, inputs[0], req[0], outputs[0]);
     auto fn = SoftmaxCompute<cpu, mxnet_op::softmax_fwd>;
diff --git a/tests/python/unittest/test_operator.py b/tests/python/unittest/test_operator.py
index f287c191963..67426693436 100644
--- a/tests/python/unittest/test_operator.py
+++ b/tests/python/unittest/test_operator.py
@@ -4098,7 +4098,7 @@ def test_new_softmax():
     for ndim in range(1, 5):
         for _ in range(5):
             shape = np.random.randint(1, 5, size=ndim)
-            axis = np.random.randint(0, ndim)
+            axis = np.random.randint(-ndim, ndim)
             data = np.random.uniform(-2, 2, size=shape)
             sym = mx.sym.softmax(axis=axis)
             check_symbolic_forward(sym, [data], [np_softmax(data, axis=axis)])


 

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services