[Op] Implement GPU version of ApplyAdamAsync, SparseApplyAdamAsync, KvSparseApplyAdamAsync.

Author: JackMoriarty

tensorflow/core/kernels/BUILD

@@ -6227,6 +6227,7 @@ tf_kernel_library(
     gpu_srcs = [
         "training_ali_ops_gpu.cu.cc",
         "training_ali_ops_gpu.h",
+        "training_ali_ops.h"
     ],
     copts = ["-g"],
     deps = [

tensorflow/core/kernels/training_ali_ops.cc

@@ -38,15 +38,30 @@ template <typename Device, typename T>
 struct ApplyAdamAsync {
   void operator()(const Device& d, typename TTypes<T>::Flat var,
                   typename TTypes<T>::Flat m, typename TTypes<T>::Flat v,
-                  typename TTypes<T>::Flat beta1_power,
-                  typename TTypes<T>::Flat beta2_power,
+                  typename TTypes<T>::Scalar beta1_power,
+                  typename TTypes<T>::Scalar beta2_power,
                   typename TTypes<T>::ConstScalar lr,
                   typename TTypes<T>::ConstScalar beta1,
                   typename TTypes<T>::ConstScalar beta2,
                   typename TTypes<T>::ConstScalar epsilon,
                   typename TTypes<T>::ConstFlat grad, bool use_nesterov);
 };
 
+template <typename Device, typename T, typename Tindex>
+struct SparseApplyAdamAsync {
+  Status operator()(const Device &d, typename TTypes<T>::Matrix var,
+                    typename TTypes<T>::Matrix m, typename TTypes<T>::Matrix v,
+                    typename TTypes<T>::Scalar beta1_power,
+                    typename TTypes<T>::Scalar beta2_power,
+                    typename TTypes<T>::ConstScalar lr,
+                    typename TTypes<T>::ConstScalar beta1,
+                    typename TTypes<T>::ConstScalar beta2,
+                    typename TTypes<T>::ConstScalar epsilon,
+                    typename TTypes<T>::ConstMatrix grad,
+                    typename TTypes<Tindex>::ConstVec indices_vec,
+                    bool apply_sparse_rmsprop, int64 inner_dim);
+};
+
 }  // end namespace functor
 }  // end namespace tensorflow