Function SparseApplyAdamKernel

tensorflow/core/kernels/training_ops_gpu.cu.cc:705–732

template <typename T, typename Tindex>
__global__ __launch_bounds__(1024) void SparseApplyAdamKernel(
  T* var, T* m, T* v, const T* grad, const T* beta1_power, const T* beta2_power,
  const T* lr, const T* beta1, const T* beta2, const T* epsilon, const Tindex* indices, 
  Tindex param_rows, Tindex updates_size, Tindex indices_size) {
    Tindex col_size = updates_size / indices_size;
    const T alpha = (*lr) * sqrt(static_cast<T>(1) - *beta2_power) /
                    (static_cast<T>(1) - *beta1_power);

    GPU_1D_KERNEL_LOOP(grad_index, updates_size) {
      Tindex indices_row = grad_index / col_size;
      Tindex param_row = indices[indices_row];
      if (param_row < 0 || param_row >= param_rows) {
        // Ignore indices that are out of range
        continue;
      }

      // Index of var, m and v
      Tindex param_index = param_row*col_size + grad_index%col_size;
      const T& g = grad[grad_index];
      T& var_a = var[param_index];
      T& m_a = m[param_index];
      T& v_a = v[param_index];

      m_a += (g-m_a) * (static_cast<T>(1) - (*beta1));
      v_a += (g*g - v_a) * (static_cast<T>(1) - (*beta2));
      var_a -= (m_a*alpha) / (sqrt(v_a)+ (*epsilon));
    }
}
template <typename T, typename Tindex> 
struct SparseApplyAdam<GPUDevice, T, Tindex> {
  Status operator()(const GPUDevice& d, typename TTypes<T>::Matrix var,