[GitHub] [incubator-mxnet] bgawrych commented on a change in pull request #20853: Improve MaskedSoftmax by oneDNN

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

bgawrych commented on a change in pull request #20853:
URL: https://github.com/apache/incubator-mxnet/pull/20853#discussion_r798556815



##########
File path: src/operator/nn/dnnl/dnnl_softmax.cc
##########
@@ -195,6 +218,24 @@ DNNLSoftmaxBwd& DNNLSoftmaxBwd::GetCached(const SoftmaxParam& param, const Tenso
   return it->second;
 }
 
+DNNLSoftmaxBwd::DNNLSoftmaxBwd(const SoftmaxParam& param, const Tensors& tensors) {
+  const float temperature   = param.temperature.has_value() ? param.temperature.value() : 1.0f;
+  const int axis            = CheckAxis(param.axis, tensors.out.shape().ndim());
+  const auto out_grad_mem   = tensors.out_grad.GetDNNLData();
+  const auto out_mem        = tensors.out.GetDNNLData();
+  const auto softmax_fwd_pd = DNNLSoftmaxFwd::GetSoftmaxFwdPd(*out_mem, axis, /*is_train*/ true);
+
+  softmax_bwd_pd = std::make_shared<softmax_bwd_pd_t>(
+      GetSoftmaxBwdPd(*out_grad_mem, *out_mem, axis, softmax_fwd_pd));
+  softmax_bwd = std::make_shared<softmax_bwd_t>(*softmax_bwd_pd);
+
+  if (temperature != 1.0f) {  // avoid dividing by 1

Review comment:
       There is no restriction in norml version so I belive we shouldn't check it either




-- 
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.

To unsubscribe, e-mail: commits-unsubscribe@mxnet.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org