```# -*- coding: utf-8 -*-
import os
from typing import Iterable, Tuple, Union

from ..tensor import Parameter, tensor
from .optimizer import Optimizer

r"""Implements AdamW algorithm proposed in `"Decoupled Weight Decay Regularization" <https://arxiv.org/abs/1711.05101>`_.

Args:
params: iterable of parameters to optimize or dicts defining
parameter groups.
lr: learning rate.
betas: coefficients used for computing running averages of gradient
and its square. Default: (0.9, 0.999)
eps: term added to the denominator to improve numerical stability. Default: 1e-8
weight_decay: weight decay (L2 penalty). Default: 1e-2
"""

def __init__(
self,
params: Union[Iterable[Parameter], dict],
lr: float,
betas: Tuple[float, float] = (0.9, 0.999),
eps: float = 1e-8,
weight_decay: float = 1e-2,
):
if lr < 0.0:
raise ValueError("Invalid learning rate: {}".format(lr))
if weight_decay < 0.0:
raise ValueError("Invalid weight_decay value: {}".format(weight_decay))
if not 0.0 <= betas[0] < 1.0:
raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
if not 0.0 <= betas[1] < 1.0:
raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))

defaults = dict(lr=lr, weight_decay=weight_decay, betas=betas, eps=eps)
super().__init__(params, defaults)
self._disable_type_convert = True

def _create_state(self, param_group):
for param in param_group["params"]:

lr = param_group["lr"]
weight_decay = param_group["weight_decay"]
eps = param_group["eps"]
beta0, beta1 = param_group["betas"]

def make_scalar(val):
return tensor(val, dtype="float32")

# since `conver_inputs` is disabled for param updates,
# scalar should be explicitly tansforred to tensor

_lr, _neg_lr = map(make_scalar, (lr, -lr))
_weight_decay = make_scalar(weight_decay)
_eps = make_scalar(eps)
_beta0, _beta1 = map(make_scalar, (beta0, beta1))

c1, c05 = map(make_scalar, (1.0, 0.5))

inplace_mode = int(os.getenv("MEGENGINE_INPLACE_UPDATE", "0"))
if inplace_mode:
# reduce device sync
c1_sub_beta0, c1_sub_beta1 = map(make_scalar, (1 - beta0, 1 - beta1))

for param in param_group["params"]:

continue

states = self._state[param]

step, exp_avg, exp_avg_sq = (
states["step"],
states["exp_avg"],
states["exp_avg_sq"],
)

if inplace_mode:
)

delta = (exp_avg / (c1 - _beta0 ** step)) / (
(exp_avg_sq / (c1 - _beta1 ** step)) ** c05 + _eps
)
if weight_decay != 0.0:
delta += param * _weight_decay
continue

# step = step + c1
step += c1

# exp_avg = _beta0 * exp_avg + grad * (c1 - _beta0)
exp_avg *= _beta0
exp_avg += grad * (c1 - _beta0)

# exp_avg_sq = _beta1 * exp_avg_sq + (c1 - _beta1) * (grad * grad)
exp_avg_sq *= _beta1