r"""Implementation of the ADOPT optimizer.

Paper: https://arxiv.org/abs/2411.02853
Original code: https://github.com/iShohei220/adopt

ADOPT (Adam with Optimal Rate for any \beta_2) is a gradient-based optimizer
designed to fix the convergence of Adam while introducing an adaptive
clipping step to stabilize updates and handle large gradients effectively.

This module provides both the high-level :class:`ADOPT` class, an `Optimizer`
subclass suitable for direct usage, and the low-level functional API
:func:`adopt` that performs the same computations. Typically, users will
instantiate the :class:`ADOPT` class.
"""

import math
from collections.abc import Callable
from typing import Any, ClassVar, cast

import torch
from composer.utils import dist
from llmfoundry.registry import optimizers
from torch import Tensor
from torch.optim.optimizer import (
    Optimizer,
    ParamsT,
    _default_to_fused_or_foreach,  # noqa: PLC2701
    _device_dtype_check_for_fused,  # noqa: PLC2701
    _disable_dynamo_if_unsupported,  # noqa: PLC2701
    _get_capturable_supported_devices,  # noqa: PLC2701
    _get_scalar_dtype,  # noqa: PLC2701
    _get_value,  # noqa: PLC2701
    _use_grad_for_differentiable,  # noqa: PLC2701
    _view_as_real,  # noqa: PLC2701
)
from torch.types import Number

from mt_dao.optimizers.utils import get_report_curvature

__all__ = ["ADOPT", "adopt"]


def _default_clip_lambda(step: Number | Tensor) -> float:
    r"""Return the default clipping value for the given step.

    By default, it uses :math:`\text{clip} = \text{step}^{0.25}`.

    Parameters
    ----------
    step : Number | Tensor | Any
        The current iteration/step value
            (can be a scalar Python number or a PyTorch Tensor).

    Returns
    -------
    float
        The clipping threshold for the given step.

    Raises
    ------
    TypeError
        If the step is not a Tensor or a Number as expected.

    """
    internal_step: int
    if isinstance(step, (Tensor)):
        internal_step = int(step.item())
    elif isinstance(step, (Number)):
        internal_step = int(step)
    else:
        msg = f"step must be a Number or Tensor, but got {type(step)}"
        raise TypeError(msg)
    return internal_step**0.25


@optimizers.register_class("adopt")
class ADOPT(Optimizer):
    r"""Implements the ADOPT algorithm.

    This is a variant of Adam designed to
    converge optimally for any :math:`\beta_2`,
    featuring an adaptive clipping
    mechanism.

    :class:`ADOPT` builds upon Adam with an additional per-step clipping of
    normalized gradients. The clipping threshold can be controlled via a user-
    defined function ``clip_lambda(step)`` or the default
    :func:`_default_clip_lambda`.

    The default threshold is :math:`\text{step}^{0.25}`, which was found to
    stabilize updates over the training process.

    This optimizer shares many of Adam's hyperparameters but modifies
    the second moment computations and update rules to incorporate this
    gradient clipping logic.

    Examples:
    --------
        >>> import torch
        >>> from adopt_optimizer import ADOPT
        >>>
        >>> model = MyModel()
        >>> optimizer = ADOPT(model.parameters(), lr=1e-3)
        >>>
        >>> for input, target in data:
        ...     def closure():
        ...         optimizer.zero_grad()
        ...         loss = loss_fn(model(input), target)
        ...         loss.backward()
        ...         return loss
        ...     loss = optimizer.step(closure)

    Notes:
    -----
        - This optimizer does not support sparse gradients.
        - The default clipping threshold scales as :math:`\text{step}^{0.25}`,
        but you can override this by passing a custom ``clip_lambda(step)``
        function.

    Warns:
    -----
        Passing a ``Tensor`` for ``lr`` is only supported if
        ``foreach=False`` or ``capturable=True``. Attempting to combine
        ``lr`` as a ``Tensor`` with ``foreach=True`` and
        ``capturable=False`` will raise an error.

    """

    metric_functions: ClassVar = {
        "l2_norm/moment": (
            lambda param, optim_state, step_tensor: torch.linalg.vector_norm(  # noqa: ARG005
                optim_state["exp_avg"],
            )
        ),
        "l2_norm/moment2": (
            lambda param, optim_state, step_tensor: torch.linalg.vector_norm(  # noqa: ARG005
                optim_state["exp_avg_sq"],
            )
        ),
        "min/param": lambda param, optim_state, step_tensor: torch.min(param.data),  # noqa: ARG005
        "max/param": lambda param, optim_state, step_tensor: torch.max(param.data),  # noqa: ARG005
        "min/moment": lambda param, optim_state, step_tensor: torch.min(  # noqa: ARG005
            optim_state["exp_avg"],
        ),
        "max/moment": lambda param, optim_state, step_tensor: torch.max(  # noqa: ARG005
            optim_state["exp_avg"],
        ),
        "min/moment2": lambda param, optim_state, step_tensor: torch.min(  # noqa: ARG005
            optim_state["exp_avg_sq"],
        ),
        "max/moment2": lambda param, optim_state, step_tensor: torch.max(  # noqa: ARG005
            optim_state["exp_avg_sq"],
        ),
        "l2_norm/param": (
            lambda param, optim_state, step_tensor: torch.linalg.vector_norm(param.data)  # noqa: ARG005
        ),
        "l2_norm/update": (
            lambda param, optim_state, step_tensor: torch.linalg.vector_norm(  # noqa: ARG005
                step_tensor,
            )
        ),
        "l2_norm/grad": (
            lambda param, optim_state, step_tensor: torch.linalg.vector_norm(param.grad)  # noqa: ARG005
        ),
    }

    def __init__(  # noqa: PLR0913, PLR0917, C901
        self,
        params: ParamsT,
        lr: float | Tensor = 1e-3,
        betas: tuple[float, float] = (0.9, 0.9999),
        eps: float = 1e-6,
        clip_lambda: (
            Callable[[Number | Tensor | Any], float] | None
        ) = _default_clip_lambda,
        weight_decay: float = 0.0,
        *,
        decouple: bool = False,
        foreach: bool | None = None,
        maximize: bool = False,
        capturable: bool = False,
        differentiable: bool = False,
        fused: bool | None = None,
        report_curvature: bool = False,
    ) -> None:
        r"""Initialise the ADOPT optimizer.

        Parameters
        ----------
        params : ParamsT
            Iterable of parameters to optimize or dicts defining parameter groups.
        lr : Union[float, Tensor], optional
            Learning rate, by default 1e-3. If a ``Tensor`` is passed, ensure you
            set ``capturable=True`` and ``foreach=False`` if you want to capture it
            in a graph.,
        betas : tuple[float, float], optional
            Coefficients used for computing running averages of gradient and
            its square, by default (0.9, 0.9999)
        eps : float, optional
            Term added to the denominator to improve numerical stability, by default
            1e-6 clip_lambda : Optional[ Callable[[Number  |  Tensor  |  Any], float] ],
            optional
            A function that, given the current
            step (as a scalar), returns the clipping threshold, by default
              _default_clip_lambda
        clip_lambda : Optional[ Callable[[Number  |  Tensor  |  Any], float] ], optional
            A function that, given the current
            step (as a scalar), returns the clipping threshold. If
            ``None``, by default _default_clip_lambda.
        weight_decay : float, optional
            Weight decay coefficient, by default 0.0.
            If ``decouple=True``, it uses a decoupled weight decay
            (sometimes known as AdamW-style). Otherwise, it follows the
            original L2 regularization approach.
        decouple : bool, optional
            Whether to decouple weight decay (like AdamW)
            or not, by default False
        foreach : Optional[bool], optional
            Whether to use the multi-tensor APIs.
            When None, it will automatically decide, by default None
        maximize : bool, optional
            Wheter to maximize rather than minimize
            the loss, by default False
        capturable : bool, optional
            Whether the optimizer should allow
            gradient capture (for example, with CUDA graphs), by default False
        differentiable : bool, optional
            Whether the optimizer is being used
            in a context that requires higher-order differentiation, by default False
        fused : Optional[bool], optional
            Whether to use a fused version of the kernel
            (if supported by PyTorch and if no constraints like
            differentiability exist), by default None
        report_curvature : bool, optional
            Whether to report curvature metrics
            for each parameter, by default False

        Raises
        ------
        ValueError
            lr as a Tensor is not supported for capturable=False and foreach=True
        RuntimeError
            `fused` is not currently supported

        """
        if not lr >= 0.0:
            msg = f"Invalid learning rate: {lr}"
            raise ValueError(msg)
        if isinstance(lr, Tensor) and foreach and not capturable:
            msg = "Tensor lr not supported for capturable=False and foreach=True"
            raise ValueError(
                msg,
            )
        if not eps >= 0.0:
            msg = f"Invalid epsilon value: {eps}"
            raise ValueError(msg)
        if not 0.0 <= betas[0] < 1.0:
            msg = f"Invalid beta parameter at index 0: {betas[0]}"
            raise ValueError(msg)
        if not 0.0 <= betas[1] < 1.0:
            msg = f"Invalid beta parameter at index 1: {betas[1]}"
            raise ValueError(msg)
        if not weight_decay >= 0.0:
            msg = f"Invalid weight_decay value: {weight_decay}"
            raise ValueError(msg)

        self.clip_lambda = clip_lambda

        defaults = {
            "lr": lr,
            "betas": betas,
            "eps": eps,
            "weight_decay": weight_decay,
            "decouple": decouple,
            "maximize": maximize,
            "foreach": foreach,
            "capturable": capturable,
            "differentiable": differentiable,
            "fused": fused,
        }
        super().__init__(params, defaults)
        for group in self.param_groups:
            group["initial_lr"] = group["lr"]

        if fused:
            if differentiable:
                msg = "`fused` does not support `differentiable`"
                raise RuntimeError(msg)
            self._step_supports_amp_scaling = True
            if foreach:
                msg = "`fused` and `foreach` cannot be `True` together."
                raise RuntimeError(msg)
            # NOTE: support fused
            msg = "`fused` is not currently supported"
            raise RuntimeError(msg)

        # NOTE: Added to avoid expensive metrics
        # calculations
        self.curvature_metric_function: (
            Callable[[Tensor, str], dict[str, Tensor]] | None
        ) = None
        if report_curvature:
            self.curvature_metric_function = get_report_curvature()

    def __setstate__(self, state: dict) -> None:
        """Set the state of the optimizer for backward compatibility.

        Specifically, this handles the presence or absence of certain keys
        in older checkpoints, ensures new keys are properly initialized,
        and converts 'step' values to Tensor if needed.

        Parameters
        ----------
        state : dict
            The optimizer state to set.

        """
        super().__setstate__(state)
        for group in self.param_groups:
            group.setdefault("maximize", False)
            group.setdefault("foreach", None)
            group.setdefault("capturable", False)
            group.setdefault("differentiable", False)
            group.setdefault("decouple", False)
            fused = group.setdefault("fused", None)
            for p in group["params"]:
                p_state = self.state.get(p, [])
                if len(p_state) != 0 and not torch.is_tensor(p_state["step"]):  # type: ignore[reportGeneralTypeIssues]
                    step_val = float(p_state["step"])  # type: ignore[reportGeneralTypeIssues]
                    p_state["step"] = (  # type: ignore[reportGeneralTypeIssues]
                        torch.tensor(
                            step_val,
                            dtype=_get_scalar_dtype(is_fused=fused),
                            device=p.device,
                        )
                        if group["capturable"] or group["fused"]
                        else torch.tensor(step_val, dtype=_get_scalar_dtype())
                    )

    def _init_group(  # noqa: PLR0913, PLR0917
        self,
        group: dict,
        params_with_grad: list[Tensor],
        grads: list[Tensor],
        exp_avgs: list[Tensor],
        exp_avg_sqs: list[Tensor],
        state_steps: list[Tensor],
    ) -> bool:
        """Initialize a parameter group for ADOPT updates.

        This internal helper function:
        - Filters out parameters that do not have gradients.
        - Checks for sparse gradients (raising an error if found).
        - Creates state entries (`step`, `exp_avg`, `exp_avg_sq`) if not already
          present.
        - Appends valid parameters and their states to the respective lists.


        Parameters
        ----------
        group : dict
            A dictionary representing a parameter group and its settings.ion_
        params_with_grad : list[Tensor]
            List to store parameters that have gradients.
        grads : list[Tensor]
            List to store gradients of those parameters.
        exp_avgs : list[Tensor]
            List to store exponential moving averages of gradients.
        exp_avg_sqs : list[Tensor]
            List to store exponential moving averages of squared gradients.
        state_steps : list[Tensor]
            List to store the optimizer step counts for these parameters.

        Returns
        -------
        bool
            A flag indicating whether complex parameters exist in this group.

        Raises
        ------
        RuntimeError
            ADOPT does not support sparse gradients
        RuntimeError
            `requires_grad` is not supported for `step` in differentiable mode
        RuntimeError
            lr as a Tensor is not supported for capturable=False and foreach=True

        """
        has_complex = False
        for p in group["params"]:
            if p.grad is None:
                continue
            has_complex |= torch.is_complex(p)
            params_with_grad.append(p)
            if p.grad.is_sparse:
                msg = "ADOPT does not support sparse gradients"
                raise RuntimeError(msg)
            grads.append(p.grad)

            state = self.state[p]

            # Lazy state initialization
            if len(state) == 0:
                if group["fused"]:
                    _device_dtype_check_for_fused(p)
                # note(crcrpar): Deliberately host `step` on CPU if both capturable and
                # fused are off. This is because kernel launches are costly on CUDA
                # and XLA.
                state["step"] = (
                    torch.zeros(
                        (),
                        dtype=_get_scalar_dtype(is_fused=group["fused"]),
                        device=p.device,
                    )
                    if group["capturable"] or group["fused"]
                    else torch.tensor(0.0, dtype=_get_scalar_dtype())
                )
                # Exponential moving average of gradient values
                state["exp_avg"] = torch.zeros_like(
                    p,
                    memory_format=torch.preserve_format,
                )
                # Exponential moving average of squared gradient values
                state["exp_avg_sq"] = torch.zeros_like(
                    p,
                    memory_format=torch.preserve_format,
                )
                # NOTE: We don't have anything related to AMSGrad here

            exp_avgs.append(state["exp_avg"])
            exp_avg_sqs.append(state["exp_avg_sq"])

            # NOTE: We don't have anything related to AMSGrad here

            if group["differentiable"] and state["step"].requires_grad:
                msg = "`requires_grad` not supported for `step` in differentiable mode"
                raise RuntimeError(
                    msg,
                )

            # Foreach without capturable does not support a tensor lr
            if (
                group["foreach"]
                and isinstance(group["lr"], Tensor)
                and not group["capturable"]
            ):
                msg = "Tensor lr is not supported for capturable=False and foreach=True"
                raise RuntimeError(
                    msg,
                )

            state_steps.append(state["step"])
        return has_complex

    @_use_grad_for_differentiable
    def step(  # type: ignore  # noqa: PGH003
        self,
        closure: Callable[[], torch.Tensor] | None = None,
    ) -> Tensor | None:
        """Perform a single ADOPT optimization step.

        Parameters
        ----------
        closure : Optional[Callable[[], torch.Tensor]], optional
            A closure that reevaluates the model
            and returns the loss. This is useful for, e.g., line search or
            other techniques that re-evaluate the model multiple times per
            iteration., by default None

        Returns
        -------
        Any | None
            The loss, if the closure is provided and called,
            otherwise None.

        """
        self._cuda_graph_capture_health_check()

        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()

        for group in self.param_groups:
            params_with_grad: list[Tensor] = []
            grads: list[Tensor] = []
            exp_avgs: list[Tensor] = []
            exp_avg_sqs: list[Tensor] = []
            state_steps: list[Tensor] = []
            beta1, beta2 = cast("tuple[float, float]", group["betas"])
            # NOTE: We don't have anything related to AMSGrad here

            has_complex = self._init_group(
                group,
                params_with_grad,
                grads,
                exp_avgs,
                exp_avg_sqs,
                state_steps,
            )

            adopt(
                params_with_grad,
                grads,
                exp_avgs,
                exp_avg_sqs,
                state_steps,
                initial_lr=group["initial_lr"],
                decouple=group["decouple"],
                clip_lambda=self.clip_lambda,
                beta1=beta1,
                beta2=beta2,
                lr=group["lr"],
                weight_decay=group["weight_decay"],
                eps=group["eps"],
                maximize=group["maximize"],
                foreach=group["foreach"],
                capturable=group["capturable"],
                differentiable=group["differentiable"],
                fused=group["fused"],
                grad_scale=getattr(self, "grad_scale", None),
                found_inf=getattr(self, "found_inf", None),
                has_complex=has_complex,
            )

        return loss

    def dist_reduce_metrics(self, optimizer_metrics: dict) -> dict:  # noqa: PLR6301
        """Reduce metrics across workers.

        Parameters
        ----------
        optimizer_metrics : dict
            The optimizer metrics per worker.

        Returns
        -------
        dict
            The reduced optimizer metrics.

        """
        local_keys = list(optimizer_metrics.keys())
        all_gathered_keys = dist.all_gather_object(local_keys)
        all_keys = set()
        for keys in all_gathered_keys:
            all_keys.update(keys)

        # Sort keys to ensure every rank has the same keys order
        # Only L2 norm metric keys are present, can apply regular sort
        all_keys = sorted(all_keys)  # type: ignore[assignment]
        for metric in all_keys:
            if metric.startswith("l2_norm"):
                reduced = optimizer_metrics.get(
                    metric,
                    torch.tensor(0.0, device=torch.cuda.current_device()),
                )
                if dist.get_world_size() > 1:
                    dist.all_reduce(reduced, reduce_operation="SUM")

                optimizer_metrics[metric] = math.sqrt(reduced)
            elif metric.startswith("min"):
                reduced = optimizer_metrics.get(
                    metric,
                    torch.tensor(0.0, device=torch.cuda.current_device()),
                )
                if dist.get_world_size() > 1:
                    dist.all_reduce(reduced, reduce_operation="MIN")

                optimizer_metrics[metric] = reduced
            elif metric.startswith("max"):
                reduced = optimizer_metrics.get(
                    metric,
                    torch.tensor(0.0, device=torch.cuda.current_device()),
                )
                if dist.get_world_size() > 1:
                    dist.all_reduce(reduced, reduce_operation="MAX")

                optimizer_metrics[metric] = reduced
            else:
                reduced = optimizer_metrics.get(
                    metric,
                    torch.tensor(0.0, device=torch.cuda.current_device()),
                )
                if dist.get_world_size() > 1:
                    dist.all_reduce(reduced, reduce_operation="SUM")
                optimizer_metrics[metric] = reduced / dist.get_world_size()

        return optimizer_metrics

    def pre_reduce_metrics(self, optimizer_metrics: dict) -> dict:  # noqa: PLR6301
        """Preprocess metrics to reduce across ranks correctly.

        Parameters
        ----------
        optimizer_metrics : dict
            The optimizer metrics, containing only the L2 metric keys.

        Returns
        -------
        dict
            The squared L2 norm metrics

        """
        # NOTE: Only L2 norm metric keys are present, can skip sorting at this stage
        for metric in optimizer_metrics:
            # NOTE: L2 norms need to be squared, before they are reduced via summation
            optimizer_metrics[metric] **= 2

        return optimizer_metrics

    def report_per_parameter_metrics(
        self,
        param: torch.Tensor,
        name: str,
        optimizer_metrics: dict,
    ) -> dict:
        """Report the metrics per parameter.

        Parameters
        ----------
        param : torch.Tensor
            The parameter for which to report metrics.
        name : str
            The name of the parameter.
        optimizer_metrics : dict
            The optimizer metrics.

        Returns
        -------
        dict
            The updated optimzier metrics.

        """
        lr = self.param_groups[0]["lr"]
        weight_decay = self.param_groups[0]["weight_decay"]
        initial_lr = self.param_groups[0]["initial_lr"]
        decouple = self.param_groups[0]["decouple"]

        if param in self.state:
            param_optim_state = self.state[param]
            # NOTE: This is inverting the ADOPT update to recover the step tensor
            step_tensor = lr * param_optim_state["exp_avg"]
            if weight_decay != 0 and decouple:
                decay_factor = (lr / initial_lr) if initial_lr else 1.0
                scaling_factor = (decay_factor * weight_decay) / (
                    1 - decay_factor * weight_decay
                )
                step_tensor.mul_(1 + scaling_factor).add_(param, alpha=scaling_factor)
            for metric in self.metric_functions:
                optimizer_metrics[f"{metric}/{name}"] = self.metric_functions[metric](
                    param,
                    param_optim_state,
                    step_tensor,
                )
            # NOTE: these are heavy and require extra memory
            if self.curvature_metric_function is not None:
                optimizer_metrics.update(self.curvature_metric_function(param, name))

        return optimizer_metrics


def _single_tensor_adopt(  # noqa: PLR0917, PLR0913
    params: list[Tensor],
    grads: list[Tensor],
    exp_avgs: list[Tensor],
    exp_avg_sqs: list[Tensor],
    state_steps: list[Tensor],
    grad_scale: Tensor | None,
    found_inf: Tensor | None,
    *,
    initial_lr: float | None,
    decouple: bool,
    clip_lambda: Callable[[Number | Tensor | Any], float] | None,
    beta1: float,
    beta2: float,
    lr: float | Tensor,
    weight_decay: float,
    eps: float,
    maximize: bool,
    capturable: bool,
    differentiable: bool,
    has_complex: bool,  # noqa: ARG001
) -> None:
    """Single-tensor implementation of ADOPT step.

    Parameters
    ----------
    params : list[Tensor]
        List of parameters to be updated.
    grads : list[Tensor]
        List of gradients corresponding to the parameters.
    exp_avgs : list[Tensor]
        List of first moment (EMA of grads).
    exp_avg_sqs : list[Tensor]
        List of second moment (EMA of squared grads).
    state_steps : list[Tensor]
        List of step counters for each parameter.
    grad_scale : Optional[Tensor]
       Not used in this implementation.
    found_inf : Optional[Tensor]
        Not used in this implementation.
    initial_lr : Optional[float]
        The initial learning rate to compute the ratio for decoupled weight decay.
    decouple : bool
        Whether to apply decoupled weight decay (AdamW style).
    clip_lambda : Optional[Callable[[Number  |  Tensor  |  Any], float]]
        Function returning the clipping threshold given the step.
    beta1 : float
        Exponential decay rate for the first moment.
    beta2 : float
        Exponential decay rate for the second moment.
    lr : Union[float, Tensor]
        Learning rate.
    weight_decay : float
        Weight decay coefficient.
    eps : float
        Term added to the denominator to improve numerical stability.
    maximize : bool
        Whether to maximize rather than minimize.
    capturable : bool
        Whether this optimizer should allow graph capture.
    differentiable : bool
        Whether this optimizer is used in a context requiring differentiable
        optimization.
    has_complex : bool
        Indicates if any parameters are complex.

    """
    assert grad_scale is None
    assert found_inf is None

    if torch.jit.is_scripting():  # type: ignore  # noqa: PGH003
        # this assert is due to JIT being dumb and not realizing that the ops below
        # have overloads to handle both float and Tensor lrs, so we just assert it's
        # a float since most people using JIT are using floats
        assert isinstance(lr, float)

    for i, param in enumerate(params):
        grad = grads[i] if not maximize else -grads[i]
        exp_avg = exp_avgs[i]
        exp_avg_sq = exp_avg_sqs[i]
        step_t = state_steps[i]

        # If compiling, the compiler will handle cudagraph checks, see note
        # [torch.compile x capturable]
        if not torch._utils.is_compiling() and capturable:  # type: ignore[reportGeneralTypeIssues]  # noqa: SLF001
            device = _get_capturable_supported_devices()
            assert param.device.type == step_t.device.type
            assert param.device.type in device, (
                f"If capturable=True, params, state_steps must be on: {device}."
            )

        step = step_t if capturable or differentiable else _get_value(step_t)

        # Perform stepweight decay if not decoupled
        if weight_decay != 0 and not decouple:
            # NOTE: The decay factor follows the same schedule of the learning rate w/o
            # being scaled by it
            decay_factor = (lr / initial_lr) if initial_lr else 1.0
            grad = grad.add(param, alpha=weight_decay)

        if torch.is_complex(param):
            grad = torch.view_as_real(grad)
            exp_avg = torch.view_as_real(exp_avg)
            exp_avg_sq = torch.view_as_real(exp_avg_sq)
            param = torch.view_as_real(param)  # noqa: PLW2901

        # NOTE: During the zero-th step (when `step == 0``), the algorithm doesn't
        # update the weights but just initializes the second momenta vectors
        if step == 0:
            # NOTE: Compared to the original implementation, we removed a useless
            # `.conj()` call to the second `grad` in the function below
            exp_avg_sq.addcmul_(grad, grad)
            # update step
            step_t += 1
            # Here, we skip whatever happens below moving to the next iteration of the
            # for loop
            continue

        # Perform stepweight decay if decoupled
        if weight_decay != 0 and decouple:
            # NOTE: The decay factor follows the same schedule of the learning rate
            # w/o being scaled by it
            decay_factor = (lr / initial_lr) if initial_lr else 1.0
            param.mul_(1 - decay_factor * weight_decay)

        denom = torch.clamp(exp_avg_sq.sqrt(), eps)
        normed_grad = grad.div(denom)
        if clip_lambda is not None:
            clip = clip_lambda(step)
            normed_grad.clamp_(-clip, clip)

        exp_avg.lerp_(normed_grad, 1 - beta1)

        param.add_(exp_avg, alpha=-_get_value(lr))
        # NOTE: Compared to the original implementation, we removed a useless `.conj()`
        # call to the second `grad` in the function below
        exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)

        # update step
        step_t += 1


def _multi_tensor_adopt(  # noqa: C901, PLR0917, PLR0913, PLR0915, PLR0912
    params: list[Tensor],
    grads: list[Tensor],
    exp_avgs: list[Tensor],
    exp_avg_sqs: list[Tensor],
    state_steps: list[Tensor],
    grad_scale: Tensor | None,
    found_inf: Tensor | None,
    *,
    initial_lr: float | None,
    has_complex: bool,
    beta1: float,
    beta2: float,
    lr: float | Tensor,
    clip_lambda: Callable[[Number | Tensor | Any], float] | None,
    weight_decay: float,
    decouple: bool,
    eps: float,
    maximize: bool,
    capturable: bool,
    differentiable: bool,
) -> None:
    """Multi-tensor implementation of the ADOPT step.

    This performs updates on groups of tensors to improve
    performance when many parameters are involved.

    Parameters
    ----------
    params : list[Tensor]
        List of parameters to be updated.
    grads : list[Tensor]
        List of gradients corresponding to the parameters.
    exp_avgs : list[Tensor]
        List of first moment (EMA of grads).
    exp_avg_sqs : list[Tensor]
        List of second moment (EMA of squared grads).
    state_steps : list[Tensor]
        List of step counters for each parameter.
    grad_scale : Optional[Tensor]
       Not used in this implementation.
    found_inf : Optional[Tensor]
        Not used in this implementation.
    initial_lr : Optional[float]
        The initial learning rate to compute the ratio for decoupled weight decay.
    decouple : bool
        Whether to apply decoupled weight decay (AdamW style).
    clip_lambda : Optional[Callable[[Number  |  Tensor  |  Any], float]]
        Function returning the clipping threshold given the step.
    beta1 : float
        Exponential decay rate for the first moment.
    beta2 : float
        Exponential decay rate for the second moment.
    lr : Union[float, Tensor]
        Learning rate.
    weight_decay : float
        Weight decay coefficient.
    eps : float
        Term added to the denominator to improve numerical stability.
    maximize : bool
        Whether to maximize rather than minimize.
    capturable : bool
        Whether this optimizer should allow graph capture.
    differentiable : bool
        Whether this optimizer is used in a context requiring differentiable
        optimization.
    has_complex : bool
        Indicates if any parameters are complex.

    Raises
    ------
    RuntimeError
        If ``lr`` is a ``Tensor`` combined with
        ``capturable=False`` and ``foreach=True``, or if the JIT compiler
        usage is not supported.

    """
    if len(params) == 0:
        return

    if isinstance(lr, Tensor) and not capturable:
        msg = "lr as a Tensor is not supported for capturable=False and foreach=True"
        raise RuntimeError(
            msg,
        )

    # If compiling, the compiler will handle cudagraph checks,
    # see note [torch.compile x capturable]
    if not torch._utils.is_compiling() and capturable:  # type: ignore[reportGeneralTypeIssues] # noqa: SLF001
        device = _get_capturable_supported_devices(
            supports_xla=False,
        )
        assert all(
            p.device.type == step.device.type and p.device.type in device
            for p, step in zip(params, state_steps, strict=False)
        ), f"If capturable=True, params, state_steps must be on: {device}."

    assert not differentiable, "_foreach ops don't support autograd"

    assert grad_scale is None
    assert found_inf is None

    grouped_tensors = Optimizer._group_tensors_by_device_and_dtype(  # noqa: SLF001
        [params, grads, exp_avgs, exp_avg_sqs, state_steps],  # type: ignore[reportArgumentType]
    )
    for (
        device_params_,
        device_grads_,
        device_exp_avgs_,
        device_exp_avg_sqs_,
        device_state_steps_,
    ), _ in grouped_tensors.values():
        device_params = cast("list[Tensor]", device_params_)
        device_grads = cast("list[Tensor]", device_grads_)
        device_exp_avgs = cast("list[Tensor]", device_exp_avgs_)
        device_exp_avg_sqs = cast("list[Tensor]", device_exp_avg_sqs_)
        device_state_steps = cast("list[Tensor]", device_state_steps_)

        # Handle complex parameters
        if has_complex:
            _view_as_real(
                device_params,
                device_grads,
                device_exp_avgs,
                device_exp_avg_sqs,
            )

        if maximize:
            device_grads = torch._foreach_neg(device_grads)  # type: ignore[assignment] # noqa: SLF001

        if weight_decay != 0 and not decouple:
            # NOTE: The decay factor follows the same schedule of the learning rate w/o
            # being scaled by it
            decay_factor = (lr / initial_lr) if initial_lr else 1.0
            weight_decay_unscaled = decay_factor * weight_decay
            # Re-use the intermediate memory (device_grads) already allocated for
            # maximize
            if maximize:
                torch._foreach_add_(
                    device_grads,
                    device_params,
                    alpha=_get_value(weight_decay_unscaled),
                )
            else:
                device_grads = torch._foreach_add(  # type: ignore[assignment] # noqa: SLF001
                    device_grads,
                    device_params,
                    alpha=-_get_value(weight_decay_unscaled),
                )

        # NOTE: During the zero-th step (when `device_state_steps[0] == 0``),
        # the algorithm doesn't update the weights but just initializes the
        # second momenta vectors
        if device_state_steps[0] == 0:
            torch._foreach_addcmul_(device_exp_avg_sqs, device_grads, device_grads)

            # NOTE: Update steps
            # If steps are on CPU, foreach will fall back to the slow path, which is a
            # for-loop calling t.add(1) over and over. 1 will then be wrapped into a
            # Tensor over and over again, which is slower than if we just wrapped it
            #  once now. The alpha is required to assure we go to the right overload.
            if (
                not torch._utils.is_compiling()  # type: ignore[reportGeneralTypeIssues] # noqa: SLF001
                and device_state_steps[0].is_cpu
            ):
                torch._foreach_add_(
                    device_state_steps,
                    torch.tensor(1.0, device="cpu"),
                    alpha=1.0,
                )
            else:
                torch._foreach_add_(device_state_steps, 1)
            # Here, we skip whatever happens below moving to the next iteration
            # of the for loop
            continue

        if weight_decay != 0 and decouple:
            # NOTE: The decay factor follows the same schedule of the learning rate
            # w/o being scaled by it
            decay_factor = (lr / initial_lr) if initial_lr else 1.0
            weight_decay_unscaled = decay_factor * weight_decay
            torch._foreach_add_(
                device_params,
                device_params,
                alpha=-_get_value(weight_decay_unscaled),
            )

        exp_avg_sq_sqrt = torch._foreach_sqrt(device_exp_avg_sqs)  # noqa: SLF001
        torch._foreach_maximum_(exp_avg_sq_sqrt, eps)

        normed_grad = torch._foreach_div(device_grads, exp_avg_sq_sqrt)  # noqa: SLF001
        if clip_lambda is not None:
            clip = clip_lambda(_get_value(device_state_steps[0]))
            torch._foreach_maximum_(normed_grad, -clip)
            torch._foreach_minimum_(normed_grad, clip)

        torch._foreach_lerp_(device_exp_avgs, normed_grad, 1 - beta1)

        torch._foreach_add_(device_params, device_exp_avgs, alpha=-_get_value(lr))
        torch._foreach_mul_(device_exp_avg_sqs, beta2)
        torch._foreach_addcmul_(
            device_exp_avg_sqs,
            device_grads,
            device_grads,
            value=1 - beta2,
        )

        # NOTE: Update steps
        # If steps are on CPU, foreach will fall back to the slow path, which is a
        # for-loop calling t.add(1) over and over. 1 will then be wrapped into a
        # Tensor over and over again, which is slower than if we just wrapped it
        # once now. The alpha is required to assure we go to the right overload.
        if not torch._utils.is_compiling() and device_state_steps[0].is_cpu:  # type: ignore[reportGeneralTypeIssues] # noqa: SLF001
            torch._foreach_add_(
                device_state_steps,
                torch.tensor(1.0, device="cpu"),
                alpha=1.0,
            )
        else:
            torch._foreach_add_(device_state_steps, 1)


def _fused_adopt(  # noqa: PLR0917,PLR0913
    params: list[Tensor],
    grads: list[Tensor],
    exp_avgs: list[Tensor],
    exp_avg_sqs: list[Tensor],
    state_steps: list[Tensor],
    grad_scale: Tensor | None,
    found_inf: Tensor | None,
    *,
    initial_lr: float | None,
    has_complex: bool,  # Needed for consistency.
    beta1: float,
    beta2: float,
    lr: float | Tensor,
    clip_lambda: Callable[[int], float] | None,
    weight_decay: float,
    decouple: bool,
    eps: float,
    maximize: bool,
    capturable: bool,  # Needed for consistency.
    differentiable: bool,
) -> None:
    """Fused implementation of the ADOPT step.

    This is not yet implemented.

    Parameters
    ----------
    params : list[Tensor]
        List of parameters to be updated.
    grads : list[Tensor]
        List of gradients corresponding to the parameters.
    exp_avgs : list[Tensor]
        List of first moment (EMA of grads).
    exp_avg_sqs : list[Tensor]
        List of second moment (EMA of squared grads).
    state_steps : list[Tensor]
        List of step counters for each parameter.
    grad_scale : Optional[Tensor]
       Not used in this implementation.
    found_inf : Optional[Tensor]
        Not used in this implementation.
    initial_lr : Optional[float]
        The initial learning rate to compute the ratio for decoupled weight decay.
    decouple : bool
        Whether to apply decoupled weight decay (AdamW style).
    clip_lambda : Optional[Callable[[Number  |  Tensor  |  Any], float]]
        Function returning the clipping threshold given the step.
    beta1 : float
        Exponential decay rate for the first moment.
    beta2 : float
        Exponential decay rate for the second moment.
    lr : Union[float, Tensor]
        Learning rate.
    weight_decay : float
        Weight decay coefficient.
    eps : float
        Term added to the denominator to improve numerical stability.
    maximize : bool
        Whether to maximize rather than minimize.
    capturable : bool
        Whether this optimizer should allow graph capture.
    differentiable : bool
        Whether this optimizer is used in a context requiring differentiable
        optimization.
    has_complex : bool
        Indicates if any parameters are complex.

    Raises
    ------
    NotImplementedError
        This function is not implemented.

    """
    raise NotImplementedError


@_disable_dynamo_if_unsupported(single_tensor_fn=_single_tensor_adopt)
def adopt(  # noqa: PLR0917,PLR0913
    params: list[Tensor],
    grads: list[Tensor],
    exp_avgs: list[Tensor],
    exp_avg_sqs: list[Tensor],
    state_steps: list[Tensor],
    # kwonly args with defaults are not supported by functions compiled with
    # torchscript issue #70627 setting this as kwarg for now as functional
    # API is compiled by torch/distributed/optim
    initial_lr: float | None = None,
    *,
    foreach: bool | None = None,
    capturable: bool = False,
    differentiable: bool = False,
    fused: bool | None = None,
    grad_scale: Tensor | None = None,
    found_inf: Tensor | None = None,
    has_complex: bool = False,
    beta1: float,
    beta2: float,
    lr: float | Tensor,
    clip_lambda: Callable[[Number | Tensor | Any], float] | None,
    weight_decay: float,
    decouple: bool,
    eps: float,
    maximize: bool,
) -> None:
    r"""Functional API that performs the ADOPT algorithm step.

    It directly applies the parameter updates in-place, modifying:
    - ``params`` in place.
    - ``exp_avgs`` (first moment).
    - ``exp_avg_sqs`` (second moment).
    - ``state_steps`` (step counters).

    It provides options to use different internal implementations:
    - Single-tensor (default).
    - Multi-tensor (``foreach=True``).
    - Fused (``fused=True``, if supported).

    The fused implementation is currently unimplemented and raises an error.

    Parameters
    ----------
    params : list[Tensor]
        List of parameters to be updated.
    grads : list[Tensor]
        List of gradients corresponding to the parameters.
    exp_avgs : list[Tensor]
        List of first moment (EMA of grads).
    exp_avg_sqs : list[Tensor]
        List of second moment (EMA of squared grads).
    state_steps : list[Tensor]
        List of step counters for each parameter.
    grad_scale : Optional[Tensor]
       Not used in this implementation.
    found_inf : Optional[Tensor]
        Not used in this implementation.
    initial_lr : Optional[float]
        The initial learning rate to compute the ratio for decoupled weight decay.
    decouple : bool
        Whether to apply decoupled weight decay (AdamW style).
    clip_lambda : Optional[Callable[[Number  |  Tensor  |  Any], float]]
        Function returning the clipping threshold given the step.
    beta1 : float
        Exponential decay rate for the first moment.
    beta2 : float
        Exponential decay rate for the second moment.
    lr : Union[float, Tensor]
        Learning rate.
    weight_decay : float
        Weight decay coefficient.
    eps : float
        Term added to the denominator to improve numerical stability.
    maximize : bool
        Whether to maximize rather than minimize.
    capturable : bool
        Whether this optimizer should allow graph capture.
    differentiable : bool
        Whether this optimizer is used in a context requiring differentiable
        optimization.
    has_complex : bool
        Indicates if any parameters are complex.
    fused : bool, optional
        Whether to use a fused version, if available.
        Currently not implemented and will raise an error if True.
        Default: None.
    foreach : bool, optional
        Whether to use multi-tensor APIs.
        Defaults to None, which auto-selects.

    Raises
    ------
    RuntimeError
        If the fused implementation is requested (``fused=True``)
        but is not supported yet, or if incompatible combinations of
        ``lr``, ``capturable``, and ``foreach`` are passed.

    Examples
    --------
        >>> import torch
        >>> from adopt_optimizer import adopt, _default_clip_lambda
        >>>
        >>> # Suppose p, grad, exp_avg, exp_avg_sq, step are Tensors
        >>> p = torch.randn(3, requires_grad=True)
        >>> grad = torch.randn_like(p)
        >>> exp_avg = torch.zeros_like(p)
        >>> exp_avg_sq = torch.zeros_like(p)
        >>> step = torch.zeros((), dtype=torch.float)
        >>>
        >>> adopt(
        ...     params=[p], grads=[grad],
        ...     exp_avgs=[exp_avg], exp_avg_sqs=[exp_avg_sq],
        ...     state_steps=[step],
        ...     lr=0.001,
        ...     beta1=0.9, beta2=0.9999,
        ...     eps=1e-6,
        ...     weight_decay=0.0,
        ...     decouple=False,
        ...     clip_lambda=_default_clip_lambda,  # optional, can be None
        ...     maximize=False,
        ... )

    """
    # NOTE: Respect when the user inputs False/True for foreach or fused. We only want
    # to change the default when neither have been user-specified. Note that we default
    # to foreach and pass False to use_fused. This is not a mistake--we want to give the
    # fused impl bake-in time before making it the default, even if it is typically
    # faster.
    if fused is None and foreach is None:
        _, foreach = _default_to_fused_or_foreach(
            params,
            differentiable,
            use_fused=False,
        )
        # Do not flip on foreach for the unsupported case where lr is a Tensor and
        # capturable=False.
        if foreach and isinstance(lr, Tensor) and not capturable:
            foreach = False
    if fused is None:
        fused = False
    if foreach is None:
        foreach = False

    if foreach and torch.jit.is_scripting():  # type: ignore # noqa: PGH003
        msg = "torch.jit.script not supported with foreach optimizers"
        raise RuntimeError(msg)
    if fused and torch.jit.is_scripting():  # type: ignore # noqa: PGH003
        msg = "torch.jit.script not supported with fused optimizers"
        raise RuntimeError(msg)

    if fused and not torch.jit.is_scripting():  # type: ignore # noqa: PGH003
        func = _fused_adopt
    elif foreach and not torch.jit.is_scripting():  # type: ignore # noqa: PGH003
        func = _multi_tensor_adopt
    else:
        func = _single_tensor_adopt

    func(
        params,
        grads,
        exp_avgs,
        exp_avg_sqs,
        state_steps,
        initial_lr=initial_lr,
        decouple=decouple,
        clip_lambda=clip_lambda,
        beta1=beta1,
        beta2=beta2,
        lr=lr,
        weight_decay=weight_decay,
        eps=eps,
        maximize=maximize,
        capturable=capturable,
        differentiable=differentiable,
        grad_scale=grad_scale,
        found_inf=found_inf,
        has_complex=has_complex,
    )
