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

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

namespace at::native::mps {

void _fused_adamw_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_adamw_" + 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_adamw_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_adamw_" + 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