   }
 }
 
-/* Fused CUDA kernel for layer normalization. It computes the LayerNorm when axis=-1.
+/* Fused CUDA kernel for the forward pass of layer normalization.
+ * It computes the LayerNorm when axis=-1, i.e., contiguous reduction scenario.
  * Shape of the input tensors:
  *      in_data = (nbatch, nchannel)
  *      gamma = (nchannel,)

  return LayerNormComputeGeneral<gpu>(attrs, ctx, inputs, req, outputs);
}

/*
 *
 */
template<int LOAD_UNROLL, bool gamma_addto, int beta_addto, typename DType>
__global__ void LayerNormFusedBackwardKernel_GammaBeta(const int nbatch,
                                                       const int nchannel,
                                                       const DType* __restrict__ in_data,
                                                       const DType* __restrict__ out_grad,
                                                       const DType* __restrict__ mean_data,
                                                       const DType* __restrict__ std_data,
                                                       DType* gamma_grad,
                                                       DType* beta_grad) {
  const int nthread = blockDim.x * blockDim.y;
  const int tid = threadIdx.x + threadIdx.y * blockDim.x;
  DType local_gamma_grad[LOAD_UNROLL];

  }
}

void GetGammaBetaGradKernelParams(const int nbatch, const int nchannel,
                                  dim3* part_grad_block_dim, dim3* part_grad_grid_dim,
                                  dim3* gb_block_dim, dim3* gb_grid_dim, int* row_repeat) {
  const int npart = 16;
  *part_grad_block_dim = dim3(32, 16);
  *part_grad_grid_dim = dim3((nchannel + 32 - 1) / 32, npart);
  *row_repeat = 16;
  *gb_block_dim = dim3(32, npart);
  *gb_grid_dim = dim3((nchannel + 32 - 1) / 32);
  CheckLaunchParam(*part_grad_grid_dim, *part_grad_block_dim);
  CheckLaunchParam(*gb_grid_dim, *gb_block_dim);
}

void LayerNormGradGPUContig(const LayerNormParam param,
                            const OpContext& ctx, const std::vector<TBlob>& inputs,
                            const std::vector<OpReqType>& req,
