# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/optimizer.fused.ipynb.

# %% ../../nbs/optimizer.fused.ipynb 1
# Contains code from:
# fastai - Apache License 2.0 - Copyright (c) 2023 fast.ai

# %% ../../nbs/optimizer.fused.ipynb 4
from __future__ import annotations
from typing import Optional, Dict

from fastcore.basics import partialler

from fastai.optimizer import (Optimizer, weight_decay, l2_reg, average_grad, sgd_step, momentum_step,
                              average_sqr_grad, rms_prop_step, step_stat, adam_step, radam_step,
                              qhadam_step, larc_layer_lr, larc_step, lamb_step, Lookahead)

from .torchscript import (JitOptimizer, radam_jit_step, sgd_jit_step, rmsprop_jit_step,
                                             adam_jit_step, radam_jit_step, qhadam_jit_step, larc_jit_step,
                                             lamb_jit_step, JitLookahead, ranger_jit_step)

from .foreach import (SGDForEachOptimizer, sgd_foreach_step, AdamForEachOptimizer,
                                         adam_foreach_step, RAdamForEachOptimizer, radam_foreach_step,
                                         LambForEachOptimizer, lamb_foreach_step, RangerForEachOptimizer,
                                         ranger_foreach_step)

from ..imports import *

# %% auto 0
__all__ = ['SGD', 'sgd', 'RMSProp', 'rmsprop', 'Adam', 'adam', 'RAdam', 'radam', 'QHAdam', 'qhadam', 'Larc', 'larc', 'Lamb',
           'lamb', 'Ranger', 'ranger']

# %% ../../nbs/optimizer.fused.ipynb 9
def SGD(
    params:Listified[Tensor], # Model parameters or parameter groups
    lr:float, # Default learning rate
    mom:float=0., # Gradient moving average (β1) coefficient
    wd:float=0., # Optional weight decay (true or L2)
    decouple_wd:bool=True, # Apply true weight decay (SGDW) or L2 regularization (SGD)
    foreach:bool=False, # Use fused ForEach implementation
    jit:bool=False # Use fused TorchScript implementation
) -> Optimizer|SGDForEachOptimizer|JitOptimizer:
    "A fastai SGD/SGDW optimizer with fused ForEach and TorchScript implementations"
    if foreach:
        return SGDForEachOptimizer(params, sgd_foreach_step, lr=lr, mom=mom, wd=wd, decouple_wd=decouple_wd)
    elif jit:
        return JitOptimizer(params, sgd_jit_step, lr=lr, mom=mom, wd=wd, decouple_wd=decouple_wd)
    else:
        cbs = [weight_decay] if decouple_wd else [l2_reg]
        if mom != 0: cbs.append(average_grad)
        cbs.append(sgd_step if mom==0 else momentum_step)
        return Optimizer(params, cbs, lr=lr, mom=mom, wd=wd)

# %% ../../nbs/optimizer.fused.ipynb 10
def sgd(
    mom:float=0., # Gradient moving average (β1) coefficient
    wd:float=0., # Optional weight decay (true or L2)
    decouple_wd:bool=True, # Apply true weight decay (SGDW) or L2 regularization (SGD)
    foreach:bool=False, # Use fused ForEach implementation
    jit:bool=False # Use fused TorchScript implementation
) -> Optimizer|SGDForEachOptimizer|JitOptimizer:
    "Partial function for the SGD/SGDW optimizer with fused ForEach and TorchScript implementations"
    return partialler(SGD, mom=mom, wd=wd, decouple_wd=decouple_wd, jit=jit, foreach=foreach)

# %% ../../nbs/optimizer.fused.ipynb 13
def RMSProp(
    params:Listified[Tensor], # Model parameters or parameter groups
    lr:float, # Default learning rate
    mom:float=0., # Gradient moving average (β1) coefficient
    sqr_mom:float=0.99, # Gradient squared moving average (β2) coefficient
    eps:float=1e-8, # Added for numerical stability
    wd:float=0., # Optional weight decay (true or L2)
    decouple_wd:bool=True, # Apply true weight decay (RMSPropW) or L2 regularization (RMSProp)
    jit:bool=False # Use fused TorchScript implementation
) -> Optimizer|JitOptimizer:
    "A fastai RMSProp/RMSPropW optimizer with a fused TorchScript implementation"
    if jit:
        return JitOptimizer(params, rmsprop_jit_step, lr=lr, mom=mom, sqr_mom=sqr_mom,
                            eps=eps, wd=wd, decouple_wd=decouple_wd)
    else:
        cbs = [weight_decay] if decouple_wd else [l2_reg]
        cbs += ([average_sqr_grad] if mom==0. else [average_grad, average_sqr_grad])
        cbs.append(rms_prop_step)
        return Optimizer(params, cbs, lr=lr, mom=mom, sqr_mom=sqr_mom, wd=wd, eps=eps)

# %% ../../nbs/optimizer.fused.ipynb 14
def rmsprop(
    mom:float=0., # Gradient moving average (β1) coefficient
    sqr_mom:float=0.99, # Gradient squared moving average (β2) coefficient
    eps:float=1e-8, # Added for numerical stability
    wd:float=0., # Optional weight decay (true or L2)
    decouple_wd:bool=True, # Apply true weight decay (RMSPropW) or L2 regularization (RMSProp)
    jit:bool=False # Use fused TorchScript implementation
) -> Optimizer|JitOptimizer:
    "Partial function for the RMSProp/RMSPropW optimizer with a fused TorchScript implementation"
    return partialler(RMSProp, mom=mom, sqr_mom=sqr_mom, eps=eps, wd=wd, decouple_wd=decouple_wd, jit=jit)

# %% ../../nbs/optimizer.fused.ipynb 17
def Adam(
    params:Listified[Tensor], # Model parameters or parameter groups
    lr:float, # Default learning rate
    mom:float=0.9, # Gradient moving average (β1) coefficient
    sqr_mom:float=0.99, # Gradient squared moving average (β2) coefficient
    eps:float=1e-5, # Added for numerical stability
    wd:float=0.01, # Optional weight decay (true or L2)
    decouple_wd:bool=True, # Apply true weight decay (AdamW) or L2 regularization (Adam)
    foreach:bool=False, # Use fused ForEach implementation
    jit:bool=False # Use fused TorchScript implementation
) -> Optimizer|AdamForEachOptimizer|JitOptimizer:
    "A fastai Adam/AdamW optimizer with fused ForEach and TorchScript implementations"
    if foreach:
        return AdamForEachOptimizer(params, adam_foreach_step, lr=lr, mom=mom,
                                    sqr_mom=sqr_mom, eps=eps, wd=wd, decouple_wd=decouple_wd)
    elif jit:
        return JitOptimizer(params, adam_jit_step, lr=lr, mom=mom, sqr_mom=sqr_mom,
                            eps=eps, wd=wd, decouple_wd=decouple_wd)
    else:
        cbs = [weight_decay] if decouple_wd else [l2_reg]
        cbs += [partial(average_grad, dampening=True), average_sqr_grad, step_stat, adam_step]
        return Optimizer(params, cbs, lr=lr, mom=mom, sqr_mom=sqr_mom, eps=eps, wd=wd)

# %% ../../nbs/optimizer.fused.ipynb 18
def adam(
    mom:float=0.9, # Gradient moving average (β1) coefficient
    sqr_mom:float=0.99, # Gradient squared moving average (β2) coefficient
    eps:float=1e-5, # Added for numerical stability
    wd:float=0.01, # Optional weight decay (true or L2)
    decouple_wd:bool=True, # Apply true weight decay (RMSPropW) or L2 regularization (RMSProp)
    foreach:bool=False, # Use fused ForEach implementation
    jit:bool=False # Use fused TorchScript implementation
) -> Optimizer|AdamForEachOptimizer|JitOptimizer:
    "Partial function for the Adam/AdamW optimizer with fused ForEach and TorchScript implementations"
    return partialler(Adam, mom=mom, sqr_mom=sqr_mom, eps=eps, wd=wd,
                      decouple_wd=decouple_wd, foreach=foreach, jit=jit)

# %% ../../nbs/optimizer.fused.ipynb 21
def RAdam(
    params:Listified[Tensor], # Model parameters or parameter groups
    lr:float, # Default learning rate
    mom:float=0.9, # Gradient moving average (β1) coefficient
    sqr_mom:float=0.99, # Gradient squared moving average (β2) coefficient
    eps:float=1e-5, # Added for numerical stability
    wd:float=0., # Optional weight decay (true or L2)
    beta:float=0., # Set to enable SAdam with native fastai RAdam
    decouple_wd:bool=True, # Apply true weight decay (RAdamW) or L2 regularization (RAdam)
    foreach:bool=False, # Use fused ForEach implementation
    jit:bool=False # Use fused TorchScript implementation
) -> Optimizer|RAdamForEachOptimizer|JitOptimizer:
    "A fastai RAdam/RAdamW optimizer with fused ForEach and TorchScript implementations"
    if (foreach or jit) and beta != 0:
        raise ValueError(f'ForEach and TorchScript RAdam does not use {beta=}, set `jit` & `foreach` to False if beta!=0')
    if foreach:
        return RAdamForEachOptimizer(params, radam_foreach_step, lr=lr, mom=mom, sqr_mom=sqr_mom,
                                     eps=eps, wd=wd, decouple_wd=decouple_wd)
    if jit:
        return JitOptimizer(params, radam_jit_step, lr=lr, mom=mom, sqr_mom=sqr_mom, eps=eps,
                            wd=wd, decouple_wd=decouple_wd)
    else:
        cbs = [weight_decay] if decouple_wd else [l2_reg]
        cbs += [partial(average_grad, dampening=True), average_sqr_grad, step_stat, radam_step]
        return Optimizer(params, cbs, lr=lr, mom=mom, sqr_mom=sqr_mom, eps=eps, wd=wd, beta=beta)

# %% ../../nbs/optimizer.fused.ipynb 22
def radam(
    mom:float=0.9, # Gradient moving average (β1) coefficient
    sqr_mom:float=0.99, # Gradient squared moving average (β2) coefficient
    eps:float=1e-5, # Added for numerical stability
    wd:float=0., # Optional weight decay (true or L2)
    beta:float=0., # Set to enable SAdam with native fastai RAdam
    decouple_wd:bool=True, # Apply true weight decay (RMSPropW) or L2 regularization (RMSProp)
    foreach:bool=False, # Use fused ForEach implementation
    jit:bool=False # Use fused TorchScript implementation
) -> Optimizer|RAdamForEachOptimizer|JitOptimizer:
    "Partial function for the RAdam/RAdamW optimizer with fused ForEach and TorchScript implementations"
    return partialler(RAdam, mom=mom, sqr_mom=sqr_mom, eps=eps, wd=wd, beta=beta,
                      decouple_wd=decouple_wd, foreach=foreach, jit=jit)

# %% ../../nbs/optimizer.fused.ipynb 25
def QHAdam(
    params:Listified[Tensor], # Model parameters or parameter groups
    lr:float, # Default learning rate
    mom:float=0.999, # Gradient moving average (β1) coefficient
    sqr_mom:float=0.999, # Gradient squared moving average (β2) coefficient
    nu_1:float=0.7, # QH immediate discount factor
    nu_2:float=1.0, # QH momentum discount factor
    eps:float=1e-8, # Added for numerical stability
    wd:float=0., # Optional weight decay (true or L2)
    decouple_wd:bool=True, # Apply true weight decay (QHAdamW) or L2 regularization (QHAdam)
    jit:bool=False # Use fused TorchScript implementation
) -> Optimizer|JitOptimizer:
    "A fastai QHAdam/QHAdamW optimizer with a fused TorchScript implementation"
    if jit:
        return JitOptimizer(params, qhadam_jit_step, lr=lr, nu_1=nu_1, nu_2=nu_2, mom=mom,
                            sqr_mom=sqr_mom, eps=eps, wd=wd, decouple_wd=decouple_wd)
    else:
        cbs = [weight_decay] if decouple_wd else [l2_reg]
        cbs += [partial(average_grad, dampening=True), average_sqr_grad, step_stat, qhadam_step]
        return Optimizer(params, cbs, lr=lr, nu_1=nu_1, nu_2=nu_2, mom=mom, sqr_mom=sqr_mom, eps=eps, wd=wd)

# %% ../../nbs/optimizer.fused.ipynb 26
def qhadam(
    mom:float=0.999, # Gradient moving average (β1) coefficient
    sqr_mom:float=0.999, # Gradient squared moving average (β2) coefficient
    nu_1:float=0.7, # QH immediate discount factor
    nu_2:float=1.0, # QH momentum discount factor
    eps:float=1e-8, # Added for numerical stability
    wd:float=0., # Optional weight decay (true or L2)
    decouple_wd:bool=True, # Apply true weight decay (RMSPropW) or L2 regularization (RMSProp)
    jit:bool=False # Use fused TorchScript implementation
) -> Optimizer|JitOptimizer:
    "Partial function for the QHAdam/QHAdamW optimizer with a fused TorchScript implementation"
    return partialler(QHAdam, mom=mom, sqr_mom=sqr_mom, nu_1=nu_1, nu_2=nu_2, eps=eps,
                      wd=wd, decouple_wd=decouple_wd, jit=jit)

# %% ../../nbs/optimizer.fused.ipynb 29
def Larc(
    params:Listified[Tensor], # Model parameters or parameter groups
    lr:float, # Default learning rate
    mom:float=0.9, # Gradient moving average (β1) coefficient
    clip:bool=True, # LARC if clip=True, LARS if clip=False
    trust_coeff:float=0.02, # Trust coeffiecnet for calculating layerwise LR
    eps:float=1e-8, # Added for numerical stability
    wd:float=0., # Optional weight decay (true or L2)
    decouple_wd:bool=True, # Apply true weight decay or L2 regularization
    jit:bool=False # Use fused TorchScript implementation
) -> Optimizer|JitOptimizer:
    "A fastai LARC/LARS optimizer with a fused TorchScript implementation"
    if jit:
        cb = partial(larc_jit_step, clip=clip)
        return JitOptimizer(params, cb, lr=lr, mom=mom, trust_coeff=trust_coeff,
                            eps=eps, wd=wd, decouple_wd=decouple_wd)
    else:
        cbs = [weight_decay] if decouple_wd else [l2_reg]
        if mom!=0.: cbs.append(average_grad)
        cbs += [partial(larc_layer_lr, clip=clip), larc_step]
        return Optimizer(params, cbs, lr=lr, mom=mom, trust_coeff=trust_coeff, eps=eps, wd=wd)

# %% ../../nbs/optimizer.fused.ipynb 30
def larc(
    mom:float=0.9, # Gradient moving average (β1) coefficient
    clip:bool=True, # LARC if clip=True, LARS if clip=False
    trust_coeff:float=0.02, # Trust coeffiecnet for calculating layerwise LR
    eps:float=1e-8, # Added for numerical stability
    wd:float=0., # Optional weight decay (true or L2)
    decouple_wd:bool=True, # Apply true weight decay (RMSPropW) or L2 regularization (RMSProp)
    jit:bool=False # Use fused TorchScript implementation
) -> Optimizer|JitOptimizer:
    "Partial function for the LARC/LARS optimizer with a fused TorchScript implementation"
    return partialler(Larc, mom=mom, clip=clip, eps=eps, trust_coeff=trust_coeff,
                      wd=wd, decouple_wd=decouple_wd, jit=jit)

# %% ../../nbs/optimizer.fused.ipynb 33
def Lamb(
    params:Listified[Tensor], # Model parameters or parameter groups
    lr:float, # Default learning rate
    mom:float=0.9, # Gradient moving average (β1) coefficient
    sqr_mom:float=0.99, # Gradient squared moving average (β2) coefficient
    eps:float=1e-5, # Added for numerical stability
    wd:float=0., # Optional weight decay (true or L2)
    decouple_wd:bool=True, # Apply true weight decay or L2 regularization
    foreach:bool=False, # Use fused ForEach implementation
    jit:bool=False # Use fused TorchScript implementation
) -> Optimizer|LambForEachOptimizer|JitOptimizer:
    "A fastai LAMB optimizer with fused ForEach and TorchScript implementations"
    if foreach:
        return LambForEachOptimizer(params, lamb_foreach_step, lr=lr, mom=mom, sqr_mom=sqr_mom,
                                    eps=eps, wd=wd, decouple_wd=decouple_wd)
    if jit:
        return JitOptimizer(params, lamb_jit_step, lr=lr, mom=mom, sqr_mom=sqr_mom,
                            eps=eps, wd=wd, decouple_wd=decouple_wd)
    else:
        cbs = [weight_decay] if decouple_wd else [l2_reg]
        cbs += [partial(average_grad, dampening=True), average_sqr_grad, step_stat, lamb_step]
        return Optimizer(params, cbs, lr=lr, mom=mom, sqr_mom=sqr_mom, eps=eps, wd=wd)

# %% ../../nbs/optimizer.fused.ipynb 34
def lamb(
    mom:float=0.9, # Gradient moving average (β1) coefficient
    sqr_mom:float=0.99, # Gradient squared moving average (β2) coefficient
    eps:float=1e-5, # Added for numerical stability
    wd:float=0., # Optional weight decay (true or L2)
    decouple_wd:bool=True, # Apply true weight decay (RMSPropW) or L2 regularization (RMSProp)
    foreach:bool=False, # Use fused ForEach implementation
    jit:bool=False # Use fused TorchScript implementation
) -> Optimizer|LambForEachOptimizer|JitOptimizer:
    "Partial function for the LAMB optimizer with fused ForEach and TorchScript implementations"
    return partialler(Lamb, mom=mom, sqr_mom=sqr_mom, eps=eps, wd=wd,
                      decouple_wd=decouple_wd, foreach=foreach, jit=jit)

# %% ../../nbs/optimizer.fused.ipynb 37
def Ranger(
    params:Listified[Tensor], # Model parameters or parameter groups
    lr:float, # Default learning rate
    mom:float=0.95, # Gradient moving average (β1) coefficient
    sqr_mom:float=0.99, # Gradient squared moving average (β2) coefficient
    eps:float=1e-6, # Added for numerical stability
    wd:float=0.01, # Optional weight decay (true or L2)
    k:int=6, # How often to conduct Lookahead step
    alpha:float=0.5, # Slow weight moving average coefficient
    decouple_wd:bool=True, # Apply true weight decay (RAdamW) or L2 regularization (RAdam)
    foreach:bool=False, # Use fused ForEach implementation
    jit:bool=False # Use fused TorchScript implementation
) -> Lookahead|RangerForEachOptimizer|JitLookahead:
    "Convenience method for `Lookahead` with `RAdam` fused ForEach and TorchScript implementations"
    if foreach:
        return RangerForEachOptimizer(params, ranger_foreach_step, lr=lr, mom=mom, sqr_mom=sqr_mom,
                                      eps=eps, wd=wd, decouple_wd=decouple_wd, k=k, alpha=alpha)
    elif jit:
        return JitLookahead(params, ranger_jit_step, lr=lr, mom=mom, sqr_mom=sqr_mom, eps=eps,
                            wd=wd, k=k, alpha=alpha, decouple_wd=decouple_wd)
    else:
        return Lookahead(RAdam(params, lr=lr, mom=mom, sqr_mom=sqr_mom, eps=eps, wd=wd,
                               decouple_wd=decouple_wd),
                         k=k, alpha=alpha)

# %% ../../nbs/optimizer.fused.ipynb 38
def ranger(
    mom:float=0.95, # Gradient moving average (β1) coefficient
    sqr_mom:float=0.99, # Gradient squared moving average (β2) coefficient
    eps:float=1e-6, # Added for numerical stability
    wd:float=0.01, # Optional weight decay (true or L2)
    k:int=6, # How often to conduct Lookahead step
    alpha:float=0.5, # Slow weight moving average coefficient
    decouple_wd:bool=True, # Apply true weight decay (RAdamW) or L2 regularization (RAdam)
    foreach:bool=False, # Use fused ForEach implementation
    jit:bool=False # Use fused TorchScript implementation
) -> Lookahead|RangerForEachOptimizer|JitLookahead:
    "Partial function of the onvenience method for `Lookahead` with `RAdam` fused ForEach and TorchScript implementations"
    return partialler(Ranger, mom=mom, sqr_mom=sqr_mom, eps=eps, wd=wd, k=k,
                      alpha=alpha, decouple_wd=decouple_wd, foreach=foreach, jit=jit)
