   AccT expx = std::exp(static_cast<AccT>(x));
   return math::TolerableValue<T>()(static_cast<T>(expx));
 }
+
+template <typename Tx, typename Ty = Tx>
+struct ExpAddFunctor {
+  HOSTDEVICE inline ExpAddFunctor(Tx max) : max(max) {}
+
+  HOSTDEVICE inline Ty operator()(const Tx& sum, const Tx& x) const {
+    return static_cast<Ty>(sum + std::exp(x - max));
+  }
+
+ private:
+  Tx max;
+};
 
 // log2(value)
 static inline int Log2Ceil(int value) {

  }
}

/*
  Wrapper of softmax with cross entropy hard label.
  - SwitchWarpSoftmaxForward for small size
  - cudnn function for large size
*/
template <typename T, typename LabelT, bool IgnoreIndex>
static void SoftmaxWithCrossEntropyHardLabel(

    T* softmax_data, int N, int dim, int D, const int ignore_index) {
  auto stream = ctx.stream();
  constexpr int max_dim = 320;
  if (D == 1) {
    if (dim <= max_dim) {  // small size
      const SoftmaxMode mode = SoftmaxMode::kCrossEntropy;
      SwitchWarpSoftmaxForward<T, LabelT, mode, IgnoreIndex>(
          loss_data, softmax_data, logits_data, labels_data, N, dim, dim,
          ignore_index, stream);
    } else {  // large size
      LaunchVectorizedSoftmaxForward<T, LabelT, IgnoreIndex>(
          loss_data, softmax_data, logits_data, labels_data, N, dim,
          ignore_index, stream);
    }
  } else {
    ScopedTensorDescriptor desc;
    std::vector<int> tensor_dims = {N, dim, D, 1};
