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

# %% ../../nbs/optimizer.foreach.ipynb 3
from __future__ import annotations

import numpy as np

from fastai.optimizer import Optimizer

from ..imports import *

# %% auto 0
__all__ = []

# %% ../../nbs/optimizer.foreach.ipynb 8
class ForEachOptimizer(Optimizer):
    "Base foreach optimizer class, updating `params` with `opt_step` instead of `Optimizer.cbs`"
    def __init__(self,
        params:Listified[Tensor], # Model parameters
        opt_step:Callable, # `ForEachOptimizer` optimizer step
        decouple_wd:bool=True, # Use true weight decay or L2 regularization, if applicable
        **defaults # Optimizer specific hyper parameters
    ):
        if notmax_torch('1.12'):
            warn(f'ForEach optimizers are untested on PyTorch {torch.__verson__}, recommended to use 1.12 or newer')
        super().__init__(params, [None], True, **defaults)
        self.opt_step = opt_step
        self.decouple_wd = decouple_wd

# %% ../../nbs/optimizer.foreach.ipynb 11
def sgd_foreach_step(p:list[Tensor], g:list[Tensor], grad_avg:list[Tensor|None], ones:list[Tensor|None], 
                     do_wd:np.ndarray[Any, bool], lr:float, wd:float, mom:float, decouple_wd:bool, **kwargs):

    if wd != 0:
        if decouple_wd:
            # weight_decay
            wd = np.where(do_wd, 1-lr*wd, 1.)
            torch._foreach_mul_(p, scalars=wd.tolist())
        else:
            # l2_reg
            wd = np.where(do_wd, wd, 1.)
            torch._foreach_addcdiv_(g, p, ones, scalars=wd.tolist())
            # cannot use scalers with foreach_add & multiple tensors, so divide by one with foreach_addcdiv

    if mom != 0:
        # average_grad
        torch._foreach_mul_(grad_avg, mom)
        torch._foreach_add_(grad_avg, g)

        # momentum_step
        torch._foreach_add_(p, grad_avg, alpha=-lr)
    else:
        # sgd_step
        torch._foreach_add_(p, g, alpha=-lr)

# %% ../../nbs/optimizer.foreach.ipynb 12
class SGDForEachOptimizer(ForEachOptimizer):
    "A `ForEachOptimizer` with a modified step for `sgd_foreach_step`"
    @torch.no_grad()
    def step(self, closure=None):
        if closure is not None:
            raise NotImplementedError("fastai optimizers currently do not support closure")
        for pg, hyper in zip(self.param_lists, self.hypers):
            pl, gl, grad_avg, ones, do_wd = [], [], [], [], []

            for p in pg:
                if hasattr(p, 'grad') and p.grad is not None:
                    state = self.state[p]

                    if 'setup' not in state:
                        if hyper['mom'] != 0:
                            state['grad_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)
                        if not self.decouple_wd:
                            state['ones'] = torch.ones_like(p, memory_format=torch.preserve_format)
                        state['setup'] = True

                    pl.append(p)
                    gl.append(p.grad)
                    grad_avg.append(state.get('grad_avg', None))
                    ones.append(state.get('ones', None))
                    do_wd.append(state.get('do_wd', True))

            self.opt_step(p=pl, g=gl, grad_avg=grad_avg, ones=ones, do_wd=np.array(do_wd, dtype=bool),
                          decouple_wd=self.decouple_wd, **hyper)

# %% ../../nbs/optimizer.foreach.ipynb 20
def adam_foreach_step(p:list[Tensor], g:list[Tensor], grad_avg:list[Tensor], sqr_avg:list[Tensor], ones:list[Tensor|None],
                      steps:np.ndarray[Any, int], do_wd:np.ndarray[Any, bool], lr:float, wd:float, mom:float, sqr_mom:float,
                      eps:float, decouple_wd:bool, **kwargs):

    if wd != 0:
        if decouple_wd:
            # weight_decay
            wd = np.where(do_wd, 1-lr*wd, 1.)
            torch._foreach_mul_(p, scalars=wd.tolist())
        else:
            # l2_reg
            wd = np.where(do_wd, wd, 1.)
            torch._foreach_addcdiv_(g, p, ones, scalars=wd.tolist())
            # cannot use scalers with foreach_add & multiple tensors, so divide by one with foreach_addcdiv

    # average_grad, dampening=True
    torch._foreach_mul_(grad_avg, mom)
    torch._foreach_add_(grad_avg, g, alpha=1-mom)

    # average_sqr_grad
    torch._foreach_mul_(sqr_avg, sqr_mom)
    torch._foreach_addcmul_(sqr_avg, g, g, value=1-sqr_mom)

    # adam_step
    debias1 = -lr / (1 - mom**steps)
    debias2 = np.sqrt(1 - sqr_mom**steps)

    sqr_avg_debias2 = torch._foreach_sqrt(sqr_avg)
    torch._foreach_div_(sqr_avg_debias2, debias2.tolist())
    torch._foreach_add_(sqr_avg_debias2, eps)

    torch._foreach_addcdiv_(p, grad_avg, sqr_avg_debias2, debias1.tolist())

# %% ../../nbs/optimizer.foreach.ipynb 21
class AdamForEachOptimizer(ForEachOptimizer):
    "An `ForEachOptimizer` with a modified step for `adam_foreach_step`"
    @torch.no_grad()
    def step(self, closure=None):
        if closure is not None:
            raise NotImplementedError("fastai optimizers currently do not support closure")
        for pg, hyper in zip(self.param_lists, self.hypers):
            pl, gl, grad_avg, sqr_avg, ones, steps, do_wd = [], [], [], [], [], [], []

            for p in pg:
                if hasattr(p, 'grad') and p.grad is not None:
                    state = self.state[p]

                    if 'step' not in state:
                        state['grad_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)
                        state['sqr_avg']  = torch.zeros_like(p, memory_format=torch.preserve_format)
                        if not self.decouple_wd:
                            state['ones'] = torch.ones_like(p, memory_format=torch.preserve_format)
                        state['step'] = 0

                    state['step'] += 1
                    pl.append(p)
                    gl.append(p.grad)
                    grad_avg.append(state['grad_avg'])
                    sqr_avg.append(state['sqr_avg'])
                    ones.append(state.get('ones', None))
                    steps.append(state['step'])
                    do_wd.append(state.get('do_wd', True))

            self.opt_step(p=pl, g=gl, grad_avg=grad_avg, sqr_avg=sqr_avg, ones=ones,
                          steps=np.array(steps, dtype=np.int32), do_wd=np.array(do_wd, dtype=bool),
                          decouple_wd=self.decouple_wd, **hyper)

# %% ../../nbs/optimizer.foreach.ipynb 28
def radam_foreach_step(p:list[Tensor], g:list[Tensor], grad_avg:list[Tensor], sqr_avg:list[Tensor], ones:list[Tensor],
                       steps:np.ndarray[Any, int], do_wd:np.ndarray[Any, bool], lr:float, wd:float, mom:float, sqr_mom:float,
                       eps:float, decouple_wd:bool, **kwargs):

    if wd != 0:
        if decouple_wd:
            # weight_decay
            wd = np.where(do_wd, 1-lr*wd, 1.)
            torch._foreach_mul_(p, scalars=wd.tolist())
        else:
            # l2_reg
            wd = np.where(do_wd, wd, 1.)
            torch._foreach_addcdiv_(g, p, ones, scalars=wd.tolist())
            # cannot use scalers with foreach_add & multiple tensors, so divide by one with foreach_addcdiv

    # average_grad, dampening=True
    torch._foreach_mul_(grad_avg, mom)
    torch._foreach_add_(grad_avg, g, alpha=1-mom)

    # average_sqr_grad
    torch._foreach_mul_(sqr_avg, sqr_mom)
    torch._foreach_addcmul_(sqr_avg, g, g, value=1-sqr_mom)

    # radam_step
    debias1 = -lr / (1 - mom**steps)
    debias2 = np.sqrt(1 - sqr_mom**steps)

    r_inf = 2/(1-sqr_mom) - 1
    r = r_inf - 2*steps*sqr_mom**steps/(1-sqr_mom**steps)

    rect   = np.where(r>5, debias1*np.sqrt(((r-4) * (r-2) * r_inf)/((r_inf-4)*(r_inf-2)*r), where=r>5), 0)
    unrect = np.where(r<=5, debias1, 0)

    # rectified step
    sqrt_avg_debias2 = torch._foreach_sqrt(sqr_avg)
    torch._foreach_div_(sqrt_avg_debias2, debias2.tolist())
    torch._foreach_add_(sqrt_avg_debias2, eps)
    torch._foreach_addcdiv_(p, grad_avg, sqrt_avg_debias2, scalars=rect.tolist())

    # unrectified step. cannot use scalers with foreach_add & multiple tensors, so divide by one with foreach_addcdiv
    torch._foreach_addcdiv_(p, grad_avg, ones, scalars=unrect.tolist())

# %% ../../nbs/optimizer.foreach.ipynb 29
class RAdamForEachOptimizer(ForEachOptimizer):
    "An `ForEachOptimizer` with a modified step for `radam_foreach_step`"
    @torch.no_grad()
    def step(self, closure=None):
        if closure is not None:
            raise NotImplementedError("fastai optimizers currently do not support closure")
        for pg, hyper in zip(self.param_lists, self.hypers):
            pl, gl, grad_avg, sqr_avg, steps, ones, do_wd = [], [], [], [], [], [], []

            for p in pg:
                if hasattr(p, 'grad') and p.grad is not None:
                    state = self.state[p]

                    if 'step' not in state:
                        state['grad_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)
                        state['sqr_avg']  = torch.zeros_like(p, memory_format=torch.preserve_format)
                        state['ones']     = torch.ones_like(p, memory_format=torch.preserve_format)
                        state['step']     = 0

                    state['step'] += 1
                    pl.append(p)
                    gl.append(p.grad)
                    grad_avg.append(state['grad_avg'])
                    sqr_avg.append(state['sqr_avg'])
                    ones.append(state['ones'])
                    steps.append(state['step'])
                    do_wd.append(state.get('do_wd', True))

            self.opt_step(p=pl, g=gl, grad_avg=grad_avg, sqr_avg=sqr_avg, ones=ones,
                          steps=np.array(steps, dtype=np.int32), do_wd=np.array(do_wd, dtype=bool),
                          decouple_wd=self.decouple_wd, **hyper)

# %% ../../nbs/optimizer.foreach.ipynb 35
@torch.jit.script
def lamb_jit_substep(p:Tensor, lstep:Tensor, lr:float):
    r1 = p.pow(2).mean().sqrt()
    r2 = lstep.pow(2).mean().sqrt()
    if r1 == 0 or r2 == 0:
        return -lr
    else:
        return -lr*min(r1/r2, 10.)

# %% ../../nbs/optimizer.foreach.ipynb 36
def lamb_foreach_step(p:list[Tensor], g:list[Tensor], grad_avg:list[Tensor], sqr_avg:list[Tensor], ones:list[Tensor],
                      steps:np.ndarray[Any, int], do_wd:np.ndarray[Any, bool], lr:float, wd:float, mom:float, sqr_mom:float,
                      eps:float, decouple_wd:bool, **kwargs):

    if wd != 0:
        if decouple_wd:
            # weight_decay
            wd = np.where(do_wd, 1-lr*wd, 1.)
            torch._foreach_mul_(p, scalars=wd.tolist())
        else:
            # l2_reg
            wd = np.where(do_wd, wd, 1.)
            torch._foreach_addcdiv_(g, p, ones, scalars=wd.tolist())
            # cannot use scalers with foreach_add & multiple tensors, so divide by one with foreach_addcdiv

    # average_grad, dampening=True
    torch._foreach_mul_(grad_avg, mom)
    torch._foreach_add_(grad_avg, g, alpha=1-mom)

    # average_sqr_grad
    torch._foreach_mul_(sqr_avg, sqr_mom)
    torch._foreach_addcmul_(sqr_avg, g, g, value=1-sqr_mom)

    # lamb_step
    debias1 = 1 - mom**steps
    debias2 = np.sqrt(1 - sqr_mom**steps)

    debias2 = torch._foreach_div(torch._foreach_sqrt(sqr_avg), debias2.tolist())
    torch._foreach_add_(debias2, eps)
    lstep = torch._foreach_div(grad_avg, debias1.tolist())
    torch._foreach_div_(lstep, debias2)

    # there are no implementations for foreach_pow, foreach_mean, or foreach_where/if methods
    q = [lamb_jit_substep(pi, ls, lr) for pi, ls in zip(p, lstep)]

    # cannot use scalers with foreach_add & multiple tensors, so divide by one with foreach_addcdiv
    torch._foreach_addcdiv_(p, lstep, ones, scalars=q)

# %% ../../nbs/optimizer.foreach.ipynb 37
class LambForEachOptimizer(RAdamForEachOptimizer):
    "An `ForEachOptimizer` with a modified step for `lamb_foreach_step`"

# %% ../../nbs/optimizer.foreach.ipynb 43
def ranger_foreach_step(p:list[Tensor], g:list[Tensor], grad_avg:list[Tensor], sqr_avg:list[Tensor], slow_p:list[Tensor],
                        ones:list[Tensor], steps:np.ndarray[Any, int], do_wd:np.ndarray[Any, bool], lr:float, wd:float,
                        mom:float, sqr_mom:float, eps:float, decouple_wd:bool, count:int, k:int, alpha:float, **kwargs):

    radam_foreach_step(p=p, g=g, grad_avg=grad_avg, sqr_avg=sqr_avg, ones=ones, steps=steps, do_wd=do_wd,
                       lr=lr, wd=wd, mom=mom, sqr_mom=sqr_mom, eps=eps, decouple_wd=decouple_wd)

    if count % k == 0:
        torch._foreach_add_(slow_p, torch._foreach_sub(p, slow_p), alpha=alpha)
        # there currently is no foreach_set method
        [pi.set_(slow_pi.clone()) for pi, slow_pi in zip(p, slow_p)]

# %% ../../nbs/optimizer.foreach.ipynb 44
class RangerForEachOptimizer(ForEachOptimizer):
    "An `ForEachOptimizer` with a modified `LookAhead` step for `ranger_foreach_step`"
    def __init__(self,
        params:Listified[Tensor], # Model parameters
        opt_step:Callable, # `ForEachOptimizer` optimizer step
        decouple_wd:bool=True, # Use true weight decay or L2 regularization, if applicable
        **defaults # Optimizer specific hyper parameters default values
    ):
        super().__init__(params, opt_step, decouple_wd, **defaults)
        self._init_state()

    @torch.no_grad()
    def step(self, closure=None):
        self.count += 1
        if closure is not None:
            raise NotImplementedError("fastai optimizers currently do not support closure")
        for pg, hyper in zip(self.param_lists, self.hypers):
            pl, gl, grad_avg, sqr_avg, slow_p, steps, ones, do_wd = [], [], [], [], [], [], [], []
            for p in pg:
                if hasattr(p, 'grad') and p.grad is not None:
                    state = self.state[p]

                    if 'step' not in state:
                        state['grad_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)
                        state['sqr_avg']  = torch.zeros_like(p, memory_format=torch.preserve_format)
                        state['ones']     = torch.ones_like(p, memory_format=torch.preserve_format)
                        state['slow_p']   = p.data.clone()
                        state['step']     = 0

                    state['step'] += 1
                    pl.append(p)
                    gl.append(p.grad)
                    grad_avg.append(state['grad_avg'])
                    sqr_avg.append(state['sqr_avg'])
                    slow_p.append(state['slow_p'])
                    ones.append(state['ones'])
                    steps.append(state['step'])
                    do_wd.append(state.get('do_wd', True))

            self.opt_step(p=pl, g=gl, grad_avg=grad_avg, sqr_avg=sqr_avg, slow_p=slow_p, ones=ones,
                          steps=np.array(steps, dtype=np.int32), do_wd=np.array(do_wd, dtype=bool),
                          decouple_wd=self.decouple_wd, count=self.count, **hyper)

    def clear_state(self):
        super().clear_state()
        self._init_state()

    def state_dict(self):
        state = super().state_dict()
        state.update({'count': self.count})
        return state

    def load_state_dict(self, sd):
        self.count = sd.pop('count')
        super().load_state_dict(sd)

    def _init_state(self):
        self.count = 0
