   return true;
 }
 
+template <>
+template <typename SIndex>
+bool RowWiseSparseAdagradOp<float, CUDAContext>::DoRunWithType() {
+  auto N = Input(GRAD).size();
+  if (N == 0) {
+    // empty grad, nothing to do here, not even launching the kernel
+    return true;
+  }
+  // size of the 1st dimension of the input gradient
+  auto GRAD_M = Input(GRAD).dim32(0);
+  auto GRAD_N = N / GRAD_M;
+
+  // each thread block will handle multiple rows of the input and output
+  RowWiseSparseAdagradKernel<<<
+      min(GRAD_M, CAFFE_MAXIMUM_NUM_BLOCKS),
+      CAFFE_CUDA_NUM_THREADS,
+      0,
+      context_.cuda_stream()>>>(
+      GRAD_M,
+      GRAD_N,
+      epsilon_,
+      Output(OUTPUT_PARAM)->template mutable_data<float>(),
+      Output(OUTPUT_MOMENT_1)->template mutable_data<float>(),
+      Input(INDICES).template data<SIndex>(),
+      Input(GRAD).template data<float>(),
+      Input(LR).template data<float>());
+  return true;
+}
+
 REGISTER_CUDA_OPERATOR(Adagrad, AdagradOp<float, CUDAContext>);
 REGISTER_CUDA_OPERATOR(SparseAdagrad, SparseAdagradOp<float, CUDAContext>);
+REGISTER_CUDA_OPERATOR(
+    RowWiseSparseAdagrad,
+    RowWiseSparseAdagradOp<float, CUDAContext>);
 }

  }
}

template<>
template<typename SIndex>
bool SparseAdagradOp<float, CUDAContext>::DoRunWithType()

 * limitations under the License.
 */

#include <cub/block/block_reduce.cuh>
#include "adagrad_op.h"
#include "caffe2/core/common_gpu.h"
#include "caffe2/core/context_gpu.h"
namespace caffe2 {

__global__ void AdagradUpdate(
