                         const std::vector<TBlob>& outputs) {
   using namespace mshadow;
   CHECK_EQ(inputs.size(), 3U);
-  mxnet::TShape data_shape(2);
-  mxnet::TShape mean_shape(1);
+  mxnet::TShape data_shape(2, 0);
+  mxnet::TShape mean_shape(1, 0);
   size_t in_ndim = inputs[layernorm::kData].ndim();
   data_shape[0] = mean_shape[0] = inputs[layernorm::kData].shape_.ProdShape(0, in_ndim - 1);
   data_shape[1] = inputs[layernorm::kData].shape_[in_ndim - 1];

 *      var_data = (nbatch,)
 *  It's always launched with (blockDim.x, blockDim.y) = (WARP_SIZE, blockDim.y)
 *  Also, when blockDim.y > 1, it requires shared memory that has size:
 *      sizeof(DType) * blockDim.y + sizeof(DType) * blockDim.y / 2
 */
template<typename DType>
__global__ void LayerNormFusedForwardKernelContig(const int nbatch,
                                                  const int nchannel,
                                                  const DType eps,
                                                  const DType* __restrict__ in_data,
                                                  const DType* __restrict__ gamma,
                                                  const DType* __restrict__ beta,

                                                  DType* data_grad) {
  int bid = blockIdx.x + blockIdx.y * gridDim.x;
  const int nthread = blockDim.x * blockDim.y;
  if (bid < nbatch) {
    // Shared memory with size blockDim.y * blockDim.x * sizeof(DType)
    extern __shared__ char buf[];
    int tid = threadIdx.x + threadIdx.y * blockDim.x;
    // 1. Calculate: mean(out_grad * gamma / std, axis=-1)
    //               mean(out_grad * gamma / std * (x - mean) / std, axis=-1)
    AType sum_val0 = 0;  // Stores mean(out_grad * gamma / std, axis=-1)
    AType sum_val1 = 0;  // Stores mean(out_grad * gamma / std * (x - mean) / std, axis=-1)
    AType mean = static_cast<AType>(mean_data[bid]);
    AType invstd_eps = AType(1) / static_cast<AType>(std_data[bid]);
    int l = LOAD_UNROLL * tid;
    for (; l + LOAD_UNROLL - 1 < nchannel; l += nthread * LOAD_UNROLL) {
#pragma unroll
      for (int i = 0; i < LOAD_UNROLL; ++i) {
        AType ele_og = static_cast<AType>(out_grad[bid * nchannel + l + i]);
        AType ele_x = static_cast<AType>(in_data[bid * nchannel + l + i]);
        AType ele_gamma = static_cast<AType>(gamma[l + i]);
        sum_val0 += ele_og * ele_gamma * invstd_eps;
        sum_val1 += ele_og * ele_gamma * (ele_x - mean) * invstd_eps * invstd_eps;
      }
    }
    for (; l < nchannel; ++l) {
      AType ele_og = static_cast<AType>(out_grad[bid * nchannel + l]);
      AType ele_x = static_cast<AType>(in_data[bid * nchannel + l]);
      AType ele_gamma = static_cast<AType>(gamma[l]);
      sum_val0 += ele_og * ele_gamma * invstd_eps;
      sum_val1 += ele_og * ele_gamma * (ele_x - mean) * invstd_eps * invstd_eps;
    }
    // Intra-warp reduction (all-reduce)
    for (int mask = blockDim.x / 2; mask > 0; mask >>= 1) {
      sum_val0 += warp_shfl_xor(sum_val0, mask);
      sum_val1 += warp_shfl_xor(sum_val1, mask);
    }
    // Inter-warp reduction (all-reduce)
    if (blockDim.y > 1) {
      AType* sum_val0_buf = reinterpret_cast<AType*>(buf);
      AType* sum_val1_buf =
        reinterpret_cast<AType*>(buf + blockDim.y / 2 * blockDim.x * sizeof(AType));
      for (int offset = blockDim.y / 2; offset > 0; offset >>= 1) {
        if (threadIdx.y >= offset && threadIdx.y < 2 * offset) {
          const int idx = (threadIdx.y - offset) * blockDim.x + threadIdx.x;
