Fix failing empty array (log_)softmax

apache · Jun 22, 2020 · be0035c · be0035c
1 parent 2fbec60
commit be0035c
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/src/operator/nn/mkldnn/mkldnn_log_softmax.cc b/src/operator/nn/mkldnn/mkldnn_log_softmax.cc
@@ -62,13 +62,15 @@ bool SupportMKLDNNLogSoftmax(const SoftmaxParam &param,
   const int in_dtype = data.dtype();
   const int out_dtype = output.dtype();
   const int axis = CheckAxis(param.axis, ndim);
+  const size_t array_size = data.shape().Size();
   // MKLDNN does not support temperature argument in their log_softmax function
   // now. Need update this once they start to support it.
   // Currently, MKLDNN shows bad performance when log_softmax is not performed on the last dimension
   if (param.temperature.has_value() ||
       in_dtype != mshadow::kFloat32 ||
-      in_dtype != out_dtype ||
-      axis != (ndim - 1)) {
+      in_dtype != out_dtype         ||
+      axis != (ndim - 1)            ||
+      array_size == 0) {
     return false;
   }
 

diff --git a/src/operator/nn/mkldnn/mkldnn_softmax.cc b/src/operator/nn/mkldnn/mkldnn_softmax.cc
@@ -64,13 +64,15 @@ bool SupportMKLDNNSoftmax(const SoftmaxParam &param,
   const int in_dtype = data.dtype();
   const int out_dtype = output.dtype();
   const int axis = CheckAxis(param.axis, ndim);
+  const size_t array_size = data.shape().Size();
   // MKLDNN does not support temperature argument in their softmax function
   // now. Need update this once they start to support it.
   // Currently, MKLDNN shows bad performance when softmax is not performed on the last dimension
   if (param.temperature.has_value() ||
       in_dtype != mshadow::kFloat32 ||
-      in_dtype != out_dtype ||
-      axis != (ndim - 1)) {
+      in_dtype != out_dtype         ||
+      axis != (ndim - 1)            ||
+      array_size == 0) {
     return false;
   }