template <typename scalar_t, typename other_t>

void adamForwardInplaceCPUKernel(const other_t b1,

const other_t inv_one_minus_pow_b1,

const other_t b2,

const other_t inv_one_minus_pow_b2,

const other_t eps,

const other_t eps_root,

const size_t n,

scalar_t *__restrict__ updates_ptr,

scalar_t *__restrict__ mu_ptr,

scalar_t *__restrict__ nu_ptr) {

#pragma omp parallel for num_threads(omp_get_num_procs())

for (size_t tid = 0; tid < n; ++tid) {

const scalar_t updates = updates_ptr[tid];

const scalar_t mu = mu_ptr[tid];

const scalar_t nu = nu_ptr[tid];

const scalar_t mu_out = b1 * mu + (1 - b1) * updates;

const scalar_t nu_out = b2 * nu + (1 - b2) * updates * updates;

const scalar_t updates_out =

mu_out * inv_one_minus_pow_b1 / (sqrt(nu_out * inv_one_minus_pow_b2 + eps_root) + eps);

mu_ptr[tid] = mu_out;

nu_ptr[tid] = nu_out;

updates_ptr[tid] = updates_out;

}

}

template <typename scalar_t, typename other_t>

void adamForwardMuCPUKernel(const scalar_t *__restrict__ updates_ptr,

const scalar_t *__restrict__ mu_ptr,

const other_t b1,

const size_t n,

scalar_t *__restrict__ mu_out_ptr) {

#pragma omp parallel for num_threads(omp_get_num_procs())

for (size_t tid = 0; tid < n; ++tid) {

const scalar_t updates = updates_ptr[tid];

const scalar_t mu = mu_ptr[tid];

const scalar_t mu_out = b1 * mu + (1 - b1) * updates;

mu_out_ptr[tid] = mu_out;

}

}

template <typename scalar_t, typename other_t>

void adamForwardNuCPUKernel(const scalar_t *__restrict__ updates_ptr,

const scalar_t *__restrict__ nu_ptr,

const other_t b2,

const size_t n,

scalar_t *__restrict__ nu_out_ptr) {

#pragma omp parallel for num_threads(omp_get_num_procs())

for (size_t tid = 0; tid < n; ++tid) {

const scalar_t updates = updates_ptr[tid];

const scalar_t nu = nu_ptr[tid];

const scalar_t nu_out = b2 * nu + (1 - b2) * pow(updates, 2);

nu_out_ptr[tid] = nu_out;

}

}

template <typename scalar_t, typename other_t>

void adamForwardUpdatesCPUKernel(const scalar_t *__restrict__ new_mu_ptr,

const scalar_t *__restrict__ new_nu_ptr,

const other_t inv_one_minus_pow_b1,

const other_t inv_one_minus_pow_b2,

const other_t eps,

const other_t eps_root,

const size_t n,

scalar_t *__restrict__ updates_out_ptr) {

#pragma omp parallel for num_threads(omp_get_num_procs())

for (size_t tid = 0; tid < n; ++tid) {

const scalar_t new_mu = new_mu_ptr[tid];

const scalar_t new_nu = new_nu_ptr[tid];

const scalar_t mu_hat = new_mu * inv_one_minus_pow_b1;

const scalar_t nu_hat = new_nu * inv_one_minus_pow_b2;

updates_out_ptr[tid] = mu_hat / (sqrt(nu_hat + eps_root) + eps);

}

}

template <typename scalar_t, typename other_t>

void adamBackwardMuCPUKernel(const scalar_t *__restrict__ dmu_ptr,

const other_t b1,

const size_t n,

scalar_t *__restrict__ dupdates_out_ptr,

scalar_t *__restrict__ dmu_out_ptr) {

#pragma omp parallel for num_threads(omp_get_num_procs())

for (size_t tid = 0; tid < n; ++tid) {

const scalar_t dmu = dmu_ptr[tid];

dupdates_out_ptr[tid] = (1 - b1) * dmu;

dmu_out_ptr[tid] = b1 * dmu;

}

}

template <typename scalar_t, typename other_t>

void adamBackwardNuCPUKernel(const scalar_t *__restrict__ dnu_ptr,

const scalar_t *__restrict__ updates_ptr,

const other_t b2,

const size_t n,

scalar_t *__restrict__ dupdates_out_ptr,

scalar_t *__restrict__ dnu_out_ptr) {

#pragma omp parallel for num_threads(omp_get_num_procs())

for (size_t tid = 0; tid < n; ++tid) {

const scalar_t dnu = dnu_ptr[tid];

const scalar_t updates = updates_ptr[tid];

dupdates_out_ptr[tid] = 2 * (1 - b2) * updates * dnu;

dnu_out_ptr[tid] = b2 * dnu;

}

}