#define TORCH_ASSERT_ONLY_METHOD_OPERATORS
#include <ATen/native/mps/operations/FusedAdamKernelImpl.h>

#include <ATen/Dispatch.h>
#include <ATen/native/ForeachUtils.h>
#include <ATen/native/mps/operations/FusedOptimizerOps.h>
#include <ATen/native/mps/operations/MultiTensorApply.h>
#include <vector>

namespace at::native::mps {

void _fused_adam_mps_impl_(TensorList params,
                           TensorList grads,
                           TensorList exp_avgs,
                           TensorList exp_avg_sqs,
                           TensorList state_steps,
                           const double lr,
                           const double beta1,
                           const double beta2,
                           const double weight_decay,
                           const double eps,
                           const bool maximize,
                           const std::optional<Tensor>& grad_scale,
                           const std::optional<Tensor>& found_inf) {
  std::vector<std::vector<Tensor>> tensor_lists{params.vec(), grads.vec(), exp_avgs.vec(), exp_avg_sqs.vec()};

  const auto kernel_name =
      "fused_adam_" + scalarToMetalTypeString(params[0]) + "_" + scalarToMetalTypeString(state_steps[0]);

  multi_tensor_apply_for_fused_optimizer<4, 512>(kernel_name,
                                                 tensor_lists,
                                                 state_steps,
                                                 FusedAdamEncodingFunctor(),
                                                 lr,
                                                 beta1,
                                                 beta2,
                                                 weight_decay,
                                                 eps,
                                                 maximize);
}

void _fused_adam_mps_impl_(TensorList params,
                           TensorList grads,
                           TensorList exp_avgs,
                           TensorList exp_avg_sqs,
                           TensorList state_steps,
                           const Tensor& lr,
                           const double beta1,
                           const double beta2,
                           const double weight_decay,
                           const double eps,
                           const bool maximize,
                           const std::optional<Tensor>& grad_scale,
                           const std::optional<Tensor>& found_inf) {
  std::vector<std::vector<Tensor>> tensor_lists{params.vec(), grads.vec(), exp_avgs.vec(), exp_avg_sqs.vec()};

  const auto kernel_name =
      "fused_adam_" + scalarToMetalTypeString(params[0]) + "_" + scalarToMetalTypeString(state_steps[0]);

  multi_tensor_apply_for_fused_optimizer<4, 512>(kernel_name,
                                                 tensor_lists,
                                                 state_steps,
                                                 FusedAdamEncodingFunctor(),
                                                 lr,
                                                 beta1,
                                                 beta2,
                                                 weight_decay,
                                                 eps,
                                                 maximize);
}
} // namespace at::native::mps