# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

import itertools
import numpy as np
import itertools
import mxnet as mx
import mxnet.lr_scheduler as lr_scheduler
from mxnet import gluon
import unittest
import pytest
import math
from mxnet.test_utils import *
from common import retry, xfail_when_nonstandard_decimal_separator

def test_learning_rate():
    o1 = mx.optimizer.Optimizer(learning_rate=0.01)
    o1.set_learning_rate(0.2)
    assert o1.learning_rate == 0.2

    lr_s = lr_scheduler.FactorScheduler(step=1)
    o2 = mx.optimizer.Optimizer(lr_scheduler=lr_s, learning_rate=0.3)
    assert o2.learning_rate == 0.3
    o2.lr_scheduler.base_lr = 0.4
    assert o2.learning_rate == 0.4

    lr_s = lr_scheduler.FactorScheduler(step=1, base_lr=1024)
    o3 = mx.optimizer.Optimizer(lr_scheduler=lr_s)
    assert o3.learning_rate == 1024


def test_learning_rate_expect_user_warning():
    lr_s = lr_scheduler.FactorScheduler(step=1)
    o = mx.optimizer.Optimizer(lr_scheduler=lr_s, learning_rate=0.3)

    with pytest.raises(UserWarning):
        o.set_learning_rate(0.5)


@xfail_when_nonstandard_decimal_separator
def test_sgd():
    opt1 = mx.optimizer.SGD
    opt2 = mx.optimizer.SGD
    shapes = [(3, 4, 5), (10, 4), (7,)]
    mom_options = [{}, {'momentum': 0.9}]
    cg_options = [{}, {'clip_gradient': 0.4}, {'clip_gradient': 0.5}]
    rg_options = [{}, {'rescale_grad': 0.14}, {'rescale_grad': 0.8}]
    wd_options = [{}, {'wd': 0.03}, {'wd': 0.05}, {'wd': 0.07}]
    mp_options = [{'multi_precision': False}, {'multi_precision': True}]
    agg_options = [{'aggregate_num': 0}, {'aggregate_num': 1},
                   {'aggregate_num': 4}, {'aggregate_num': np.inf}]

    for dtype in [np.float16, np.float32]:
        for params in itertools.product(mom_options, cg_options, rg_options,
                                        wd_options, mp_options, agg_options):
            kwarg = {k: v for param in params for k, v in param.items()}
            if (dtype == np.float16 and ('multi_precision' not in kwarg or
                                         not kwarg['multi_precision'])):
                continue
            if dtype == np.float16:
                compare_optimizer(opt1(use_fused_step=False, **kwarg),
                                  opt2(use_fused_step=True, **kwarg),
                                  shapes, dtype, rtol=1e-3, atol=1e-4)
            else:
                compare_optimizer(opt1(use_fused_step=False, **kwarg),
                                  opt2(use_fused_step=True, **kwarg),
                                  shapes, dtype)
            # test operator fallback on cpu
            if dtype != np.float16:
                compare_optimizer(opt1(use_fused_step=False, **kwarg),
                                  opt2(use_fused_step=True, **kwarg),
                                  [shapes[0][:2], shapes[1]],
                                  dtype, w_stype='csr', g_stype='csr')


class PySparseSGD(mx.optimizer.Optimizer):
    """python reference implemenation of sgd"""
    def __init__(self, learning_rate=0.1, momentum=0.0, **kwargs):
        super(PySparseSGD, self).__init__(learning_rate=learning_rate, **kwargs)
        self.momentum = momentum

    def create_state(self, index, weight):
        """Create additional optimizer state: momentum

        Parameters
        ----------
        weight : NDArray
        The weight data

        """
        if self.momentum == 0.0:
            return None
        else:
            return mx.nd.zeros(weight.shape, weight.context, dtype=weight.dtype)

    def step(self, indices, weights, grads, states):
        """Perform an optimization step using gradients and states.

        Parameters
        ----------
        indices : list of int
            List of unique indices of the parameters into the individual learning rates
            and weight decays. Learning rates and weight decay may be set via `set_lr_mult()`
            and `set_wd_mult()`, respectively.
        weights : list of NDArray
            List of parameters to be updated.
        grads : list of NDArray
            List of gradients of the objective with respect to this parameter.
        states : List of any obj
            List of state returned by `create_state()`.
        """
        for index, weight, grad, state in zip(indices, weights, grads, states):
            lr = self._get_lr(index)
            wd = self._get_wd(index)
            self._update_count(index)
            num_rows = weight.shape[0]
            if self.momentum == 0.0:
                # Update on a per row basis, skip all-zero rows
                for row in range(num_rows):
                    grad_row = grad[row].asnumpy()
                    all_zeros = mx.test_utils.almost_equal(grad_row, np.zeros_like(grad_row))
                    if all_zeros:
                        continue
                    grad[row] *= self.rescale_grad
                    if self.clip_gradient is not None:
                        grad[row] = mx.nd.clip(grad[row], -self.clip_gradient, self.clip_gradient)
                    grad[row] += wd * weight[row]
                    weight[row] -= lr * grad[row]
            else:
                mom = state
                for row in range(num_rows):
                    grad_row = grad[row].asnumpy()
                    all_zeros = mx.test_utils.almost_equal(grad_row, np.zeros_like(grad_row))
                    if all_zeros:
                        continue
                    grad[row] *= self.rescale_grad
                    if self.clip_gradient is not None:
                        grad[row] = mx.nd.clip(grad[row], -self.clip_gradient, self.clip_gradient)
                    grad[row] += wd * weight[row]
                    mom[row] *= self.momentum
                    mom[row] -= lr * grad[row]
                    weight[row] += mom[row]


@xfail_when_nonstandard_decimal_separator
def test_sparse_sgd():
    opt1 = PySparseSGD
    opt2 = mx.optimizer.SGD
    shapes = [(3, 4, 5), (10, 4), (7,)]
    mom_options = [{}, {'momentum': 0.9}]
    cg_options = [{}, {'clip_gradient': 0.4}, {'clip_gradient': 0.5}]
    rg_options = [{}, {'rescale_grad': 0.14}, {'rescale_grad': 0.8}]
    wd_options = [{}, {'wd': 0.03}, {'wd': 0.05}]
    agg_options = [{'aggregate_num': 0}, {'aggregate_num': 1},
                   {'aggregate_num': 4}, {'aggregate_num': np.inf}]
    for dtype in [np.float32]:
        for params in itertools.product(mom_options, cg_options, rg_options,
                                        wd_options, agg_options):
            kwarg = {k: v for param in params for k, v in param.items()}
            compare_optimizer(opt1(**kwarg),
                              opt2(use_fused_step=True, lazy_update=True, **kwarg), shapes, dtype,
                              w_stype='row_sparse', g_stype='row_sparse')
            compare_optimizer(opt1(**kwarg),
                              opt2(use_fused_step=True, lazy_update=True, **kwarg), shapes, dtype,
                              w_stype='default', g_stype='row_sparse')


@xfail_when_nonstandard_decimal_separator
def test_std_sparse_sgd():
    opt1 = mx.optimizer.SGD
    opt2 = mx.optimizer.SGD
    shapes = [(3, 4, 5), (10, 4), (7,)]
    mom_options = [{}, {'momentum': 0.9}]
    cg_options = [{}, {'clip_gradient': 0.4}, {'clip_gradient': 0.5}]
    rg_options = [{}, {'rescale_grad': 0.14}, {'rescale_grad': 0.8}]
    wd_options = [{}, {'wd': 0.03}, {'wd': 0.05}, {'wd': 0.07}]
    agg_options = [{'aggregate_num': 0}, {'aggregate_num': 1},
                   {'aggregate_num': 4}, {'aggregate_num': np.inf}]

    for dtype in [np.float32]:
        for params in itertools.product(mom_options, cg_options, rg_options,
                                        wd_options, agg_options):
            kwarg = {k: v for param in params for k, v in param.items()}
            compare_optimizer(opt1(use_fused_step=False, **kwarg),
                              opt2(use_fused_step=True, lazy_update=False, **kwarg), shapes, dtype,
                              w_stype='row_sparse', g_stype='row_sparse')
            compare_optimizer(opt1(use_fused_step=False, **kwarg),
                              opt2(use_fused_step=True, lazy_update=False, **kwarg), shapes, dtype,
                              w_stype='default', g_stype='row_sparse')


@xfail_when_nonstandard_decimal_separator
def test_nag():
    opt1 = mx.optimizer.NAG
    opt2 = mx.optimizer.NAG
    shapes = [(3, 4, 5), (10, 4), (7,)]
    mom_options = [{}, {'momentum': 0.9}]
    cg_options = [{}, {'clip_gradient': 0.4}, {'clip_gradient': 0.5}]
    rg_options = [{}, {'rescale_grad': 0.14}, {'rescale_grad': 0.8}]
    wd_options = [{}, {'wd': 0.03}, {'wd': 0.05}]
    mp_options = [{'multi_precision': False}, {'multi_precision': True}]
    agg_options = [{'aggregate_num': 0}, {'aggregate_num': 1},
                   {'aggregate_num': 4}, {'aggregate_num': np.inf}]

    for dtype in [np.float16, np.float32]:
        for params in itertools.product(mom_options, cg_options, rg_options,
                                        wd_options, mp_options, agg_options):
            kwarg = {k: v for param in params for k, v in param.items()}
            if (dtype == np.float16 and ('multi_precision' not in kwarg or
                                         not kwarg['multi_precision'])):
                continue
            compare_optimizer(opt1(use_fused_step=False, **kwarg),
                              opt2(use_fused_step=True, **kwarg),
                              shapes, dtype, rtol=1e-3, atol=1e-4)


@xfail_when_nonstandard_decimal_separator
def test_lars():
    opt1 = mx.optimizer.LARS
    opt2 = mx.optimizer.LARS
    shapes = [(3, 4, 5), (10, 4), (7,)]
    eta_options = [{}, {'eta': 0.002}, {'eta': 0.01}]
    mom_options = [{'momentum': 0.0}, {'momentum': 0.9}]
    cg_options = [{}, {'clip_gradient': 0.4}]
    rg_options = [{}, {'rescale_grad': 0.14}]
    wd_options = [{}, {'wd': 0.03}, {'wd': 0.05}]
    mp_options = [{'multi_precision': False}, {'multi_precision': True}]
    agg_options = [{'aggregate_num': 0}, {'aggregate_num': 1},
                   {'aggregate_num': 4}, {'aggregate_num': np.inf}]
    for dtype in [np.float16, np.float32]:
        for params in itertools.product(eta_options, mom_options, cg_options, rg_options,
                                        wd_options, mp_options, agg_options):
            kwarg = {k: v for param in params for k, v in param.items()}
            if (dtype == np.float16 and ('multi_precision' not in kwarg or
                                         not kwarg['multi_precision'])):
                continue
            compare_optimizer(opt1(use_fused_step=False, **kwarg),
                              opt2(use_fused_step=True, **kwarg),
                              shapes, dtype, rtol=1e-3, atol=1e-3)


@xfail_when_nonstandard_decimal_separator
def test_lamb():
    opt1 = mx.optimizer.LAMB
    opt2 = mx.optimizer.LAMB

    shapes = [(3, 4, 5), (10, 4), (7,)]
    beta1_options = [{}, {'beta1': 0.5}]
    beta2_options = [{}, {'beta2': 0.8}]
    cg_options = [{}, {'clip_gradient': 0.4}]
    rg_options = [{}, {'rescale_grad': 0.14}]
    wd_options = [{}, {'wd': 0.03}]
    bc_options = [{'bias_correction': False}, {'bias_correction': True}]
    lb_options = [{'lower_bound': None}, {'lower_bound': 1e-3}]
    ub_options = [{'upper_bound': None}, {'upper_bound': 10}]
    mp_options = [{'multi_precision': False}, {'multi_precision': True}]
    agg_options = [{'aggregate_num': 0}, {'aggregate_num': 1},
                   {'aggregate_num': 4}]
    for dtype in [np.float16, np.float32]:
        for params in itertools.product(beta1_options, beta2_options, cg_options, rg_options,
                                        wd_options, bc_options, lb_options, ub_options,
                                        mp_options, agg_options):
            kwarg = {k: v for param in params for k, v in param.items()}
            if (dtype == np.float16 and ('multi_precision' not in kwarg or
                                         not kwarg['multi_precision'])):
                continue
            compare_optimizer(opt1(use_fused_step=False, **kwarg),
                              opt2(use_fused_step=True, **kwarg),
                              shapes, dtype, rtol=1e-3, atol=1e-3)


@xfail_when_nonstandard_decimal_separator
def test_lans():
    opt1 = mx.optimizer.LANS
    opt2 = mx.optimizer.LANS

    shapes = [(3, 4, 5), (10, 4), (7,)]
    beta1_options = [{}, {'beta1': 0.5}]
    beta2_options = [{}, {'beta2': 0.8}]
    cg_options = [{}, {'clip_gradient': 0.4}]
    rg_options = [{}, {'rescale_grad': 0.14}]
    wd_options = [{}, {'wd': 0.03}]
    lb_options = [{'lower_bound': None}, {'lower_bound': 1e-3}]
    ub_options = [{'upper_bound': None}, {'upper_bound': 10}]
    mp_options = [{'multi_precision': False}, {'multi_precision': True}]
    agg_options = [{'aggregate_num': 0}, {'aggregate_num': 1},
                   {'aggregate_num': 4}]
    for dtype in [np.float16, np.float32]:
        for params in itertools.product(beta1_options, beta2_options, cg_options, rg_options,
                                        wd_options, lb_options, ub_options,
                                        mp_options, agg_options):
            kwarg = {k: v for param in params for k, v in param.items()}
            if (dtype == np.float16 and ('multi_precision' not in kwarg or
                                         not kwarg['multi_precision'])):
                continue
            compare_optimizer(opt1(use_fused_step=False, **kwarg),
                              opt2(use_fused_step=True, **kwarg),
                              shapes, dtype, rtol=1e-3, atol=1e-3)


def test_sgld():
    opt1 = mx.optimizer.SGLD
    opt2 = mx.optimizer.SGLD
    shapes = [(3, 4, 5), (10, 4), (7,)]
    ns_options = [1234, 42]
    cg_options = [{}, {'clip_gradient': 0.4}, {'clip_gradient': 0.5}]
    wd_options = [{}, {'wd': 0.03}, {'wd': 0.05}, {'wd': 0.07}]
    mp_options = [{'multi_precision': False}, {'multi_precision': True}]
    agg_options = [{'aggregate_num': 0}, {'aggregate_num': 1},
                   {'aggregate_num': 4}, {'aggregate_num': np.inf}]

    for seed in ns_options:
        for dtype in [np.float16, np.float32]:
            for params in itertools.product(cg_options, wd_options, mp_options, agg_options):
                kwarg = {k: v for param in params for k, v in param.items()}
                if (dtype == np.float16 and ('multi_precision' not in kwarg
                                             or not kwarg['multi_precision'])):
                    continue
                atol = 1e-2 if dtype == np.float16 else 1e-3
                rtol = 1e-4 if dtype == np.float16 else 1e-5
                compare_optimizer_noise_seeded(opt1(**kwarg),
                                               opt2(**kwarg),
                                               shapes, dtype, seed, atol=atol, rtol=rtol)


@xfail_when_nonstandard_decimal_separator
def test_ftml():
    opt1 = mx.optimizer.FTML
    opt2 = mx.optimizer.FTML
    shapes = [(3, 4, 5), (10, 4), (7,)]
    beta1_options = [{}, {'beta1': 0.5}, {'beta1': 0.7}]
    beta2_options = [{}, {'beta2': 0.8}, {'beta2': 0.9}]
    cg_options = [{}, {'clip_gradient': 0.4}, {'clip_gradient': 0.5}]
    rg_options = [{}, {'rescale_grad': 0.14}, {'rescale_grad': 0.8}]
    wd_options = [{}, {'wd': 0.03}, {'wd': 0.05}, {'wd': 0.07}]
    mp_options = [{'multi_precision': False}, {'multi_precision': True}]
    agg_options = [{'aggregate_num': 0}, {'aggregate_num': 1},
                   {'aggregate_num': 4}, {'aggregate_num': np.inf}]

    for dtype in [np.float16, np.float32]:
        for params in itertools.product(beta1_options, beta2_options, cg_options,
                                        rg_options, wd_options, mp_options, agg_options):
            kwarg = {k: v for param in params for k, v in param.items()}
            if (dtype == np.float16 and ('multi_precision' not in kwarg or
                                         not kwarg['multi_precision'])):
                continue
            compare_optimizer(opt1(use_fused_step=False, **kwarg),
                              opt2(use_fused_step=True, **kwarg),
                              shapes, dtype, rtol=1e-3, atol=1e-4)


# Sparse ADAM
class PySparseAdam(mx.optimizer.Optimizer):
    """python reference implemenation of sparse adam"""
    def __init__(self, learning_rate=0.001, beta1=0.9, beta2=0.999, epsilon=1e-8,
                 lazy_update=False, **kwargs):
        super(PySparseAdam, self).__init__(learning_rate=learning_rate, **kwargs)
        self.beta1 = beta1
        self.beta2 = beta2
        self.epsilon = epsilon
        self.lazy_update = lazy_update

    def create_state(self, index, weight):
        """Create additional optimizer state: mean, variance

        Parameters
        ----------
        weight : NDArray
        The weight data

        """
        return (mx.nd.zeros(weight.shape, weight.context, dtype=weight.dtype),  # mean
                mx.nd.zeros(weight.shape, weight.context, dtype=weight.dtype))  # variance

    def step(self, indices, weights, grads, states):
        """Perform an optimization step using gradients and states.

        Parameters
        ----------
        indices : list of int
            List of unique indices of the parameters into the individual learning rates
            and weight decays. Learning rates and weight decay may be set via `set_lr_mult()`
            and `set_wd_mult()`, respectively.
        weights : list of NDArray
            List of parameters to be updated.
        grads : list of NDArray
            List of gradients of the objective with respect to this parameter.
        states : List of any obj
            List of state returned by `create_state()`.
        """
        for index, weight, grad, state in zip(indices, weights, grads, states):
            self._update_count(index)
            lr = self._get_lr(index)
            wd = self._get_wd(index)
            t = self._index_update_count[index]

            mean, variance = state
            num_rows = weight.shape[0]

            coef1 = 1. - self.beta1 ** t
            coef2 = 1. - self.beta2 ** t
            lr *= math.sqrt(coef2) / coef1

            for row in range(num_rows):
                # check row slices of all zeros
                all_zeros = mx.test_utils.almost_equal(grad[row].asnumpy(),
                                                       np.zeros_like(grad[row].asnumpy()))
                # skip zeros during lazy update
                if all_zeros and self.lazy_update:
                    continue
                grad[row] *= self.rescale_grad
                # clip gradients
                if self.clip_gradient is not None:
                    mx.nd.clip(grad[row], -self.clip_gradient, self.clip_gradient, out=grad[row])
                grad[row] += wd * weight[row]
                # update mean
                mean[row] *= self.beta1
                mean[row] += grad[row] * (1. - self.beta1)
                # update variance
                variance[row] *= self.beta2
                variance[row] += (1 - self.beta2) * mx.nd.square(grad[row], out=grad[row])
                # update weight
                weight[row] -= lr * mean[row] / (mx.nd.sqrt(variance[row]) + self.epsilon)


@xfail_when_nonstandard_decimal_separator
def test_adam():
    opt1 = mx.optimizer.Adam
    opt2 = mx.optimizer.Adam
    shapes = [(3, 4, 5), (10, 4), (7,)]
    beta1_options = [{}, {'beta1': 0.5}, {'beta1': 0.7}]
    beta2_options = [{}, {'beta2': 0.8}, {'beta2': 0.9}]
    cg_options = [{}, {'clip_gradient': 0.4}, {'clip_gradient': 0.5}]
    rg_options = [{}, {'rescale_grad': 0.14}, {'rescale_grad': 0.8}]
    wd_options = [{}, {'wd': 0.03}, {'wd': 0.05}, {'wd': 0.07}]
    mp_options = [{'multi_precision': False}, {'multi_precision': True}]
    agg_options = [{'aggregate_num': 0}, {'aggregate_num': 1},
                   {'aggregate_num': 4}, {'aggregate_num': np.inf}]
    for dtype in [np.float16, np.float32]:
        for params in itertools.product(beta1_options, beta2_options, cg_options,
                                        rg_options, wd_options, mp_options, agg_options):
            kwarg = {k: v for param in params for k, v in param.items()}
            if (dtype == np.float16 and ('multi_precision' not in kwarg or
                                         not kwarg['multi_precision'])):
                continue
            # atol 2e-5 needed to pass with seed 1248389097
            compare_optimizer(opt1(use_fused_step=False, **kwarg),
                              opt2(use_fused_step=True, **kwarg), shapes, dtype,
                              rtol=1e-4, atol=2e-5)


@xfail_when_nonstandard_decimal_separator
def test_sparse_adam():
    opt1 = PySparseAdam
    opt2 = mx.optimizer.Adam
    shapes = [(3, 4, 5), (10, 4), (7,)]
    beta1_options = [{}, {'beta1': 0.5}]
    beta2_options = [{}, {'beta2': 0.8}]
    cg_options = [{}, {'clip_gradient': 0.4}]
    rg_options = [{}, {'rescale_grad': 0.14}]
    wd_options = [{}, {'wd': 0.03}]
    mp_options = [{'multi_precision': False}, {'multi_precision': True}]
    agg_options = [{'aggregate_num': 0}, {'aggregate_num': 1},
                   {'aggregate_num': 4}, {'aggregate_num': np.inf}]
    for dtype in [np.float16, np.float32]:
        for params in itertools.product(beta1_options, beta2_options, cg_options,
                                        rg_options, wd_options, mp_options, agg_options):
            kwarg = {k: v for param in params for k, v in param.items()}
            if (dtype == np.float16 and ('multi_precision' not in kwarg or
                                         not kwarg['multi_precision'])):
                continue
            # atol 2e-5 needed to pass with seed 1248389097
            compare_optimizer(opt1(lazy_update=False, **kwarg),
                              opt2(use_fused_step=True, lazy_update=False, **kwarg), shapes, dtype,
                              rtol=1e-4, atol=2e-5)
            # atol 2e-5 needed to pass with seed 781809840
            compare_optimizer(opt1(lazy_update=True, **kwarg),
                              opt2(use_fused_step=True, lazy_update=True, **kwarg), shapes,
                              dtype, w_stype='row_sparse', g_stype='row_sparse',
                              rtol=1e-4, atol=2e-5)
            compare_optimizer(opt1(lazy_update=False, **kwarg),
                              opt2(use_fused_step=True, lazy_update=False, **kwarg), shapes,
                              dtype, w_stype='row_sparse', g_stype='row_sparse',
                              rtol=1e-4, atol=2e-5)
            compare_optimizer(opt1(lazy_update=True, **kwarg),
                              opt2(use_fused_step=True, lazy_update=True, **kwarg), shapes,
                              dtype, w_stype='default', g_stype='row_sparse',
                              rtol=1e-4, atol=2e-5)
            compare_optimizer(opt1(lazy_update=False, **kwarg),
                              opt2(use_fused_step=True, lazy_update=False, **kwarg), shapes,
                              dtype, w_stype='default', g_stype='row_sparse',
                              rtol=1e-4, atol=2e-5)


@xfail_when_nonstandard_decimal_separator
@pytest.mark.skip(reason="Flaky test https://github.com/apache/incubator-mxnet/issues/18400")
def test_adamax():
    opt1 = mx.optimizer.Adamax
    opt2 = mx.optimizer.Adamax
    shapes = [(3, 4, 5), (10, 4), (7,)]
    beta1_options = [{}, {'beta1': 0.5}, {'beta1': 0.7}]
    beta2_options = [{}, {'beta2': 0.8}, {'beta2': 0.9}]
    cg_options = [{}, {'clip_gradient': 0.4}, {'clip_gradient': 0.5}]
    rg_options = [{}, {'rescale_grad': 0.14}, {'rescale_grad': 0.8}]
    wd_options = [{}, {'wd': 0.03}, {'wd': 0.05}, {'wd': 0.07}]
    mp_options = [{'multi_precision': False}, {'multi_precision': True}]
    agg_options = [{'aggregate_num': 0}, {'aggregate_num': 1},
                   {'aggregate_num': 4}, {'aggregate_num': np.inf}]
    for dtype in [np.float16, np.float32]:
        for params in itertools.product(beta1_options, beta2_options, cg_options,
                                        rg_options, wd_options, mp_options, agg_options):
            kwarg = {k: v for param in params for k, v in param.items()}
            if (dtype == np.float16 and
                    ('multi_precision' not in kwarg or not kwarg['multi_precision'])):
                continue
            compare_optimizer(opt1(**kwarg), opt2(**kwarg), shapes, dtype)


@xfail_when_nonstandard_decimal_separator
def test_signum():
    opt1 = mx.optimizer.Signum
    opt2 = mx.optimizer.Signum
    shapes = [(3, 4, 5), (10, 4), (7,)]
    cg_options = [{}, {'clip_gradient': 0.4}, {'clip_gradient': 0.5}]
    rg_options = [{}, {'rescale_grad': 0.14}, {'rescale_grad': 0.8}]
    wd_options = [{}, {'wd': 0.03}, {'wd': 0.05}, {'wd': 0.07}]
    wd_lh_options = [{}, {'wd_lh': 0.015}, {'wd_lh': 0.0}]
    mom_options = [{}, {'momentum': 0.9}]
    lr_options = [{'learning_rate': 0.05},{'learning_rate': 0.01}]
    mp_options = [{'multi_precision': False}, {'multi_precision': True}]
    agg_options = [{'aggregate_num': 0}, {'aggregate_num': 1},
                   {'aggregate_num': 4}, {'aggregate_num': np.inf}]
    for dtype in [np.float16, np.float32]:
        for params in itertools.product(cg_options, rg_options, wd_options,
                                        wd_lh_options, mom_options, lr_options,
                                        mp_options, agg_options):
            kwarg = {k: v for param in params for k, v in param.items()}
            if (dtype == np.float16 and
                    ('multi_precision' not in kwarg or not kwarg['multi_precision'])):
                continue
            rtol, atol = (1e-3, 1e-4) if dtype is np.float16 else (1e-4, 1e-5)
            compare_optimizer(opt1(use_fused_step=False, **kwarg),
                              opt2(use_fused_step=True, **kwarg), shapes, dtype,
                                   rtol=rtol, atol=atol)


@xfail_when_nonstandard_decimal_separator
def test_rms():
    opt1 = mx.optimizer.RMSProp
    opt2 = mx.optimizer.RMSProp
    shapes = [(3, 4, 5), (10, 4), (7,)]
    rho_options = [{}, {'rho': 0.5}]
    cg_options = [{}, {'clip_gradient': 0.4}]
    cw_options = [{}, {'clip_weights': 0.01}]
    center_options = [{'centered': False}, {'centered': True}]
    rg_options = [{}, {'rescale_grad': 0.14}]
    wd_options = [{}, {'wd': 0.03}]
    mom_options = [{'momentum': 0.0}, {'momentum': 0.9}]
    mp_options = [{'multi_precision': False}, {'multi_precision': True}]
    agg_options = [{'aggregate_num': 0}, {'aggregate_num': 1},
                   {'aggregate_num': 4}, {'aggregate_num': np.inf}]
    for dtype in [np.float16, np.float32]:
        # Reduce foating point compare tolerance to avoid flaky test failure.
        rtol, atol = (1e-1, 1e-1) if dtype is np.float16 else (1e-2, 1e-2)

        for params in itertools.product(rho_options, cg_options, cw_options,
                                        center_options, rg_options, wd_options,
                                        mom_options, mp_options, agg_options):
            kwarg = {k: v for param in params for k, v in param.items()}
            if (dtype == np.float16 and
                    ('multi_precision' not in kwarg or not kwarg['multi_precision'])):
                continue
            compare_optimizer(opt1(use_fused_step=False, **kwarg),
                              opt2(use_fused_step=True, **kwarg), shapes, dtype,
                              rtol=rtol, atol=atol)
            if default_device() == mx.cpu():
                compare_optimizer(opt1(use_fused_step=False, **kwarg),
                                  opt2(use_fused_step=True, **kwarg),
                                  shapes, dtype, g_stype='row_sparse', rtol=rtol, atol=atol)


class PySparseFtrl(mx.optimizer.Optimizer):
    """python reference implemenation of sparse Ftrl optimizer.

    Referenced from *Ad Click Prediction: a View from the Trenches*, available at
    http://dl.acm.org/citation.cfm?id=2488200.

    Parameters
    ----------
    lamda1 : float, optional
        L1 regularization coefficient.
    learning_rate : float, optional
        The initial learning rate.
    beta : float, optional
        Per-coordinate learning rate correlation parameter.
    eta :
        .. math::
           \\eta_{t,i} = \\frac{learningrate}{\\beta+\\sqrt{\\sum_{s=1}^tg_{s,i}^t}}
    """

    def __init__(self, lamda1=0.01, learning_rate=0.1, beta=1, **kwargs):
        super(PySparseFtrl, self).__init__(**kwargs)
        self.lamda1 = lamda1
        self.beta = beta
        self.lr = learning_rate

    def create_state(self, index, weight):
        return (mx.nd.zeros(weight.shape, weight.context, dtype=weight.dtype),  # z
                mx.nd.zeros(weight.shape, weight.context, dtype=weight.dtype))  # n

    def step(self, indices, weights, grads, states):
        """Perform an optimization step using gradients and states.

        Parameters
        ----------
        indices : list of int
            List of unique indices of the parameters into the individual learning rates
            and weight decays. Learning rates and weight decay may be set via `set_lr_mult()`
            and `set_wd_mult()`, respectively.
        weights : list of NDArray
            List of parameters to be updated.
        grads : list of NDArray
            List of gradients of the objective with respect to this parameter.
        states : List of any obj
            List of state returned by `create_state()`.
        """
        for index, weight, grad, state in zip(indices, weights, grads, states):
            self._update_count(index)
            wd = self._get_wd(index)
            lr = self._get_lr(index)
            num_rows = weight.shape[0]

            z, n = state
            for row in range(num_rows):
                all_zeros = mx.test_utils.almost_equal(grad[row].asnumpy(), np.zeros_like(grad[row].asnumpy()))
                if all_zeros:
                    continue
                grad[row] *= self.rescale_grad
                if self.clip_gradient is not None:
                    mx.nd.clip(grad[row], -self.clip_gradient, self.clip_gradient, out=grad[row])

                # update z[row], n[row]
                sigma = - mx.nd.sqrt(n[row])
                n[row] += mx.nd.square(grad[row])
                denom = mx.nd.sqrt(n[row])
                sigma += denom
                sigma /= lr
                z[row] += grad[row] - sigma * weight[row]

                # update weight
                denom += self.beta
                denom /= lr
                denom += wd
                d = mx.nd.sign(z[row]) * mx.nd.maximum(mx.nd.abs(z[row]) - self.lamda1, 0)
                weight[row] = - d / denom


@xfail_when_nonstandard_decimal_separator
@retry(3)
def test_ftrl():
    opt1 = mx.optimizer.Ftrl
    opt2 = mx.optimizer.Ftrl
    shapes = [(3, 4, 5), (10, 4), (7,)]
    lamda1_options = [{'lamda1': 0.}, {'lamda1': 0.1}]
    cg_options = [{}, {'clip_gradient': 0.4}, {'clip_gradient': 0.5}]
    rg_options = [{}, {'rescale_grad': 0.14}, {'rescale_grad': 0.8}]
    wd_options = [{}, {'wd': 0.03}, {'wd': 0.05}, {'wd': 0.07}]
    mp_options = [{'multi_precision': False}, {'multi_precision': True}]
    agg_options = [{'aggregate_num': 0}, {'aggregate_num': 1},
                   {'aggregate_num': 4}, {'aggregate_num': np.inf}]
    for dtype in [np.float16, np.float32]:
        for params in itertools.product(lamda1_options, cg_options,
                                        rg_options, wd_options,
                                        mp_options, agg_options):
            kwarg = {k: v for param in params for k, v in param.items()}
            if (dtype == np.float16 and
                    ('multi_precision' not in kwarg or not kwarg['multi_precision'])):
                continue
            rtol, atol = (1e-3, 1e-3) if dtype is np.float16 else (1e-4, 1e-4)
            compare_optimizer(opt1(use_fused_step=False, **kwarg),
                              opt2(use_fused_step=True, **kwarg), shapes, dtype,
                              rtol=rtol, atol=atol)


@xfail_when_nonstandard_decimal_separator
def test_sparse_ftrl():
    opt1 = PySparseFtrl
    opt2 = mx.optimizer.Ftrl
    shapes = [(3, 4, 5), (10, 4), (7,)]
    lamda1_options = [{'lamda1': 0.}, {'lamda1': 0.1}]
    cg_options = [{}, {'clip_gradient': 0.4}, {'clip_gradient': 0.5}]
    rg_options = [{}, {'rescale_grad': 0.14}, {'rescale_grad': 0.8}]
    wd_options = [{}, {'wd': 0.03}, {'wd': 0.05}, {'wd': 0.07}]
    mp_options = [{'multi_precision': False}, {'multi_precision': True}]
    agg_options = [{'aggregate_num': 0}, {'aggregate_num': 1},
                   {'aggregate_num': 4}, {'aggregate_num': np.inf}]
    for dtype in [np.float16, np.float32]:
        for params in itertools.product(lamda1_options, cg_options,
                                        rg_options, wd_options,
                                        mp_options, agg_options):
            kwarg = {k: v for param in params for k, v in param.items()}
            if (dtype == np.float16 and
                    ('multi_precision' not in kwarg or not kwarg['multi_precision'])):
                continue
            rtol, atol = (1e-3, 1e-3) if dtype is np.float16 else (1e-4, 1e-4)
            compare_optimizer(opt1(**kwarg), opt2(**kwarg), shapes,
                              dtype, w_stype='row_sparse', g_stype='row_sparse',
                              rtol=rtol, atol=atol)


@xfail_when_nonstandard_decimal_separator
def test_nadam():
    opt1 = mx.optimizer.Nadam
    opt2 = mx.optimizer.Nadam
    shapes = [(3, 4, 5), (10, 4), (7,)]
    beta1_options = [{}, {'beta1': 0.5}]
    beta2_options = [{}, {'beta2': 0.8}]
    schedule_decay_options = [{}, {'schedule_decay': 0.008}]
    cg_options = [{}, {'clip_gradient': 0.4}, {'clip_gradient': 0.5}]
    rg_options = [{}, {'rescale_grad': 0.14}, {'rescale_grad': 0.8}]
    wd_options = [{}, {'wd': 0.03}, {'wd': 0.05}]
    mp_options = [{'multi_precision': False}, {'multi_precision': True}]
    agg_options = [{'aggregate_num': 0}, {'aggregate_num': 1},
                   {'aggregate_num': 4}, {'aggregate_num': np.inf}]
    for dtype in [np.float16, np.float32]:
        for params in itertools.product(beta1_options, beta2_options, cg_options,
                                        schedule_decay_options, rg_options, wd_options,
                                        mp_options, agg_options):
            kwarg = {k: v for param in params for k, v in param.items()}
            if (dtype == np.float16 and
                    ('multi_precision' not in kwarg or not kwarg['multi_precision'])):
                continue
            compare_optimizer(opt1(**kwarg), opt2(**kwarg), shapes, dtype)


class PySparseAdaGrad(mx.optimizer.Optimizer):
    """python reference implemenation of sparse Adagrad optimizer.

    This class implements the AdaGrad optimizer described in *Adaptive Subgradient
    Methods for Online Learning and Stochastic Optimization*, and available at
    http://www.jmlr.org/papers/volume12/duchi11a/duchi11a.pdf.

    Parameters
    ----------
    learning_rate : float, default 0.01
        The initial learning rate. If None, the optimization will use the
        learning rate from ``lr_scheduler``. If not None, it will overwrite
        the learning rate in ``lr_scheduler``. If None and ``lr_scheduler``
        is also None, then it will be set to 0.01 by default.
    epsilon : float, default 1e-6
        Small value to avoid division by 0.
    """

    def __init__(self, learning_rate=0.01, epsilon=1e-6, **kwargs):
        super(PySparseAdaGrad, self).__init__(learning_rate=learning_rate,
                                              **kwargs)
        self.epsilon = epsilon

    def create_state(self, index, weight):
        return mx.nd.zeros(weight.shape, weight.context, stype=weight.stype)  # history

    def step(self, indices, weights, grads, states):
        """Perform an optimization step using gradients and states.

        Parameters
        ----------
        indices : list of int
            List of unique indices of the parameters into the individual learning rates
            and weight decays. Learning rates and weight decay may be set via `set_lr_mult()`
            and `set_wd_mult()`, respectively.
        weights : list of NDArray
            List of parameters to be updated.
        grads : list of NDArray
            List of gradients of the objective with respect to this parameter.
        states : List of any obj
            List of state returned by `create_state()`.
        """
        for index, weight, grad, state in zip(indices, weights, grads, states):
            self._update_count(index)
            wd = self._get_wd(index)
            lr = self._get_lr(index)
            num_rows = weight.shape[0]

            history = state
            for row in range(num_rows):
                all_zeros = mx.test_utils.almost_equal(grad[row].asnumpy(), np.zeros_like(grad[row].asnumpy()))
                if all_zeros:
                    continue
                grad[row] *= self.rescale_grad
                if self.clip_gradient is not None:
                    mx.nd.clip(grad[row], -self.clip_gradient, self.clip_gradient, out=grad[row])
                grad[row] += wd * weight[row]

                # update history[row]
                history[row] += mx.nd.square(grad[row])
                denom = mx.nd.sqrt(history[row])
                denom += self.epsilon

                # update weight
                weight[row] -= lr * grad[row] / denom


def test_adagrad():
    opt1 = mx.optimizer.AdaGrad
    opt2 = mx.optimizer.AdaGrad
    shapes = [(3, 4, 5), (10, 4), (7,)]
    eps_options = [{}, {'epsilon': 1e-8}]
    cg_options = [{}, {'clip_gradient': 0.4}, {'clip_gradient': 0.5}]
    rg_options = [{}, {'rescale_grad': 0.14}, {'rescale_grad': 0.8}]
    wd_options = [{}, {'wd': 0.0}]
    agg_options = [{'aggregate_num': 0}, {'aggregate_num': 1},
                   {'aggregate_num': 4}, {'aggregate_num': np.inf}]
    for dtype in [np.float16, np.float32]:
        for params in itertools.product(eps_options, cg_options,
                                        rg_options, wd_options, agg_options):
            kwarg = {k: v for param in params for k, v in param.items()}
            if dtype is np.float16:
                kwarg.update({'multi_precision': True})
            compare_optimizer(opt1(use_fused_step=False, **kwarg),
                              opt2(use_fused_step=True, **kwarg), shapes, dtype)


@xfail_when_nonstandard_decimal_separator
def test_sparse_adagrad():
    opt1 = PySparseAdaGrad
    opt2 = mx.optimizer.AdaGrad
    shapes = [(3, 4, 5), (10, 4), (7,)]
    eps_options = [{}, {'epsilon': 1e-8}]
    cg_options = [{}, {'clip_gradient': 0.4}, {'clip_gradient': 0.5}]
    rg_options = [{}, {'rescale_grad': 0.14}, {'rescale_grad': 0.8}]
    wd_options = [{}, {'wd': 0.0}]
    agg_options = [{'aggregate_num': 0}, {'aggregate_num': 1},
                   {'aggregate_num': 4}, {'aggregate_num': np.inf}]
    for dtype in [np.float16, np.float32]:
        for params in itertools.product(eps_options, cg_options,
                                        rg_options, wd_options, agg_options):
            kwarg = {k: v for param in params for k, v in param.items()}
            if dtype is np.float16:
                kwarg.update({'multi_precision': True})
            if kwarg.get('wd', 0.0) == 0.0:
                compare_optimizer(opt1(**kwarg), opt2(use_fused_step=True, **kwarg), shapes, dtype,
                                  w_stype='row_sparse', g_stype='row_sparse')
                compare_optimizer(opt1(**kwarg), opt2(use_fused_step=True, **kwarg), shapes, dtype,
                                  g_stype='row_sparse')


def test_adadelta():
    opt1 = mx.optimizer.AdaDelta
    opt2 = mx.optimizer.AdaDelta
    shapes = [(3, 4, 5), (10, 4), (7,)]
    rho_options = [{'rho': 0.9}]
    eps_options = [{}, {'epsilon': 1e-8}]
    cg_options = [{}, {'clip_gradient': 0.4}, {'clip_gradient': 0.5}]
    rg_options = [{}, {'rescale_grad': 0.14}, {'rescale_grad': 0.8}]
    wd_options = [{}, {'wd': 0.03}, {'wd': 0.05}, {'wd': 0.07}]
    agg_options = [{'aggregate_num': 0}, {'aggregate_num': 1},
                   {'aggregate_num': 4}, {'aggregate_num': np.inf}]
    for dtype in [np.float16, np.float32]:
        for params in itertools.product(rho_options, eps_options, cg_options,
                                        rg_options, wd_options, agg_options):
            kwarg = {k: v for param in params for k, v in param.items()}
            if dtype is np.float16:
                kwarg.update({'multi_precision': True})
            compare_optimizer(opt1(**kwarg), opt2(**kwarg), shapes, dtype)


def test_dcasgd():
    opt1 = mx.optimizer.DCASGD
    opt2 = mx.optimizer.DCASGD
    shapes = [(3, 4, 5), (10, 4), (7,)]
    lamda_options = [{}, {'lamda': 0.01}, {'lamda': 0.1}]
    mom_options = [{}, {'momentum': 0.0}, {'momentum': 0.9}]
    cg_options = [{}, {'clip_gradient': 0.4}, {'clip_gradient': 0.5}]
    rg_options = [{}, {'rescale_grad': 0.14}, {'rescale_grad': 0.8}]
    wd_options = [{}, {'wd': 0.03}, {'wd': 0.05}, {'wd': 0.07}]
    agg_options = [{'aggregate_num': 0}, {'aggregate_num': 1},
                   {'aggregate_num': 4}, {'aggregate_num': np.inf}]
    for dtype in [np.float16, np.float32]:
        for params in itertools.product(lamda_options, mom_options, cg_options,
                                        rg_options, wd_options, agg_options):
            kwarg = {k: v for param in params for k, v in param.items()}
            if dtype is np.float16:
                kwarg.update({'multi_precision': True})
            compare_optimizer(opt1(**kwarg), opt2(**kwarg), shapes, dtype)


def test_adamW():
    opt1 = mx.optimizer.AdamW
    opt2 = mx.optimizer.AdamW
    shapes = [(3, 4, 5), (10, 4), (7,)]
    beta1_options = [{}, {'beta1': 0.5}, {'beta1': 0.7}]
    beta2_options = [{}, {'beta2': 0.8}, {'beta2': 0.9}]
    cg_options = [{}, {'clip_gradient': 0.4}, {'clip_gradient': 0.5}]
    rg_options = [{}, {'rescale_grad': 0.14}, {'rescale_grad': 0.8}]
    wd_options = [{}, {'wd': 0.03}, {'wd': 0.05}, {'wd': 0.07}]
    mp_options = [{'multi_precision': False}, {'multi_precision': True}]
    agg_options = [{'aggregate_num': 0}, {'aggregate_num': 1},
                   {'aggregate_num': 4}, {'aggregate_num': np.inf}]
    correct_bias_options = [{'correct_bias': True}, {'correct_bias': False}]
    for dtype in [np.float16, np.float32]:
        for params in itertools.product(beta1_options, beta2_options, cg_options,
                                        rg_options, wd_options, mp_options,
                                        agg_options, correct_bias_options):
            kwarg = {k: v for param in params for k, v in param.items()}
            if (dtype == np.float16 and ('multi_precision' not in kwarg or
                                         not kwarg['multi_precision'])):
                continue
            compare_optimizer(opt1(use_fused_step=False, **kwarg),
                              opt2(use_fused_step=True, **kwarg), shapes, dtype,
                              rtol=1e-3, atol=2e-3)

def test_adabelief():
    opt1 = mx.optimizer.AdaBelief
    opt2 = mx.optimizer.AdaBelief
    shapes = [(3, 4, 5), (10, 4), (7,)]
    beta1_options = [{}, {'beta1': 0.5}, {'beta1': 0.7}]
    beta2_options = [{}, {'beta2': 0.8}, {'beta2': 0.9}]
    cg_options = [{}, {'clip_gradient': 0.4}, {'clip_gradient': 0.5}]
    rg_options = [{}, {'rescale_grad': 0.14}, {'rescale_grad': 0.8}]
    wd_options = [{}, {'wd': 0.03}, {'wd': 0.05}, {'wd': 0.07}]
    mp_options = [{'multi_precision': False}, {'multi_precision': True}]
    agg_options = [{'aggregate_num': 0}, {'aggregate_num': 1},
                   {'aggregate_num': 4}, {'aggregate_num': np.inf}]
    correct_bias_options = [{'correct_bias': True}, {'correct_bias': False}]
    for dtype in [np.float16, np.float32]:
        for params in itertools.product(beta1_options, beta2_options, cg_options,
                                        rg_options, wd_options, mp_options,
                                        agg_options, correct_bias_options):
            kwarg = {k: v for param in params for k, v in param.items()}
            if (dtype == np.float16 and ('multi_precision' not in kwarg or
                                         not kwarg['multi_precision'])):
                continue
            compare_optimizer(opt1(use_fused_step=False, **kwarg),
                              opt2(use_fused_step=True, **kwarg), shapes, dtype,
                              rtol=1e-3, atol=2e-3)

def test_factor_scheduler():
    base_lr = 1
    step = 100
    factor = 0.1
    sched = mx.lr_scheduler.FactorScheduler(step, factor, stop_factor_lr=1e-4, base_lr=base_lr,
                                        warmup_steps=20, warmup_begin_lr=0.1, warmup_mode='constant')

    assert (sched(0) == 0.1)
    np.testing.assert_almost_equal(sched(10), 0.1)
    assert (sched(21) == base_lr), sched(21)
    np.testing.assert_almost_equal(sched(101), base_lr * factor)
    np.testing.assert_almost_equal(sched(201), base_lr * factor * factor)
    np.testing.assert_almost_equal(sched(1000), 1e-4)


def test_multifactor_scheduler():
    base_lr = 0.1
    steps = [15, 25]
    factor = 0.1
    sched = mx.lr_scheduler.MultiFactorScheduler(steps, factor, base_lr=base_lr,
                                        warmup_steps=10, warmup_begin_lr=0.05, warmup_mode='linear')

    assert sched(0) == 0.05
    np.testing.assert_almost_equal(sched(5), 0.05 + (base_lr - 0.05)/2)
    np.testing.assert_almost_equal(sched(15), base_lr)
    np.testing.assert_almost_equal(sched(16), base_lr * factor)
    np.testing.assert_almost_equal(sched(20), base_lr * factor)
    np.testing.assert_almost_equal(sched(26), base_lr * factor * factor)
    np.testing.assert_almost_equal(sched(100), base_lr * factor * factor)


def test_poly_scheduler():
    base_lr = 3
    final_lr = 0
    steps = 1000
    poly_sched = mx.lr_scheduler.PolyScheduler(steps, base_lr=base_lr, pwr=2, final_lr=final_lr,
                                    warmup_steps=100, warmup_begin_lr=0, warmup_mode='linear')

    np.testing.assert_almost_equal(poly_sched(0), 0)
    np.testing.assert_almost_equal(poly_sched(50), float(base_lr)/2)
    np.testing.assert_almost_equal(poly_sched(100), base_lr)
    assert (poly_sched(101) <  poly_sched(100))
    assert (poly_sched(500) < 1.6)
    np.testing.assert_almost_equal(poly_sched(steps), final_lr)


def test_cosine_scheduler():
    # also tests case without warmup
    base_lr = 3
    final_lr = 0.1
    steps = 1000
    cosine_sched = mx.lr_scheduler.CosineScheduler(steps, base_lr=base_lr, final_lr=final_lr)
    np.testing.assert_almost_equal(cosine_sched(0), base_lr)
    np.testing.assert_almost_equal(cosine_sched(steps), final_lr)
    assert (cosine_sched(500) > 1.5)