 
   float* g_a_data = g_a->mutable_data<float>();
   float* g_b_data = g_b->mutable_data<float>();
-  math::Set<float, CUDAContext>(g_a->size(), 0.f, g_a_data, &context_);
-  math::Set<float, CUDAContext>(g_b->size(), 0.f, g_b_data, &context_);
 
-  ElementwiseLinearGradientKernel<<<CAFFE_GET_BLOCKS(D), CAFFE_CUDA_NUM_THREADS,
-                                  0, context_.cuda_stream()>>>(
-    N, D, g_o.data<float>(), X.data<float>(), a.data<float>(),
-    g_X->mutable_data<float>(), g_a_data, g_b_data);
+  ElementwiseLinearGradientKernel<<<
+      D,
+      CAFFE_CUDA_NUM_THREADS,
+      0,
+      context_.cuda_stream()>>>(
+      N,
+      D,
+      g_o.data<float>(),
+      X.data<float>(),
+      a.data<float>(),
+      g_X->mutable_data<float>(),
+      g_a_data,
+      g_b_data);
   return true;
 }
 

#include "caffe2/core/context_gpu.h"
#include "caffe2/operators/operator_fallback_gpu.h"


namespace caffe2 {


__global__ void ElementwiseLinearGradientKernel(const int N, const int D,
  const float* g_o_data, const float* X_data, const float* a_data,
  float* g_X_data, float* g_a_data, float* g_b_data) {
  int d = blockIdx.x; // One block per D

  float g_a_sum = 0;
  float g_b_sum = 0;
  for (int n = threadIdx.x; n < N; n += blockDim.x) {
    const float gox = g_o_data[n * D + d];
    g_X_data[n * D + d] = gox * a_data[d];
    g_a_sum += gox * X_data[n * D + d];
    g_b_sum += gox;
  }

  typedef cub::BlockReduce<float, CAFFE_CUDA_NUM_THREADS> BlockReduce;
  __shared__ typename BlockReduce::TempStorage temp_storage;

  float g_a_sum_tot = BlockReduce(temp_storage).Sum(g_a_sum);
  float g_b_sum_tot = BlockReduce(temp_storage).Sum(g_b_sum);

  if (threadIdx.x == 0) {
    g_a_data[d] = g_a_sum_tot;
    g_b_data[d] = g_b_sum_tot;
  }
}

