# 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.

"""Some of the tests using CUDNN require a special GPU instruction called dp4a.
Ref: http://images.nvidia.com/content/pdf/tesla/184457-Tesla-P4-Datasheet-NV-Final-Letter-Web.pdf
"""
import os
import mxnet as mx
import numpy as onp
from mxnet import npx
from mxnet.util import use_np
from mxnet.gluon.model_zoo import vision
from mxnet.test_utils import assert_almost_equal, assert_exception, rand_ndarray, rand_shape_nd, same, DummyIter
from common import xfail_when_nonstandard_decimal_separator
from mxnet.io import NDArrayIter
import unittest
import operator

npx.reset_np()

def collect_block_args_aux(block, sym):
  arg_params, aux_params = dict(), dict()
  for k, v in block.collect_params().items():
    if k in sym.list_arguments():
      arg_params[k]= v._reduce()
    elif k in sym.list_auxiliary_states():
      aux_params[k]= v._reduce()
  return arg_params, aux_params

def is_test_for_gpu():
    return mx.current_device().device_type == 'gpu'


def is_test_for_dnnl():
    return (mx.current_device().device_type == 'cpu'
            and os.environ.get('ENABLE_ONEDNN_QUANTIZATION_TEST') == '1')


def is_test_for_native_cpu():
    return (mx.current_device().device_type == 'cpu'
            and os.environ.get('ENABLE_ONEDNN_QUANTIZATION_TEST') == None)


def get_low_high(qtype):
    """ Return low and high value for given integer type as float number"""
    if qtype == 'uint8':
        return 0.0, 255.0
    else:
        return -128.0, 127.0


def test_quantize_float32_to_int8():
    shape = rand_shape_nd(4)
    data = rand_ndarray(shape, 'default', dtype='float32')
    min_range = mx.nd.min(data)
    max_range = mx.nd.max(data)
    qdata, min_val, max_val = mx.nd.contrib.quantize(data, min_range, max_range, out_type='int8')
    data_np = data.asnumpy()
    min_range = min_range.asscalar()
    max_range = max_range.asscalar()
    real_range = onp.maximum(onp.abs(min_range), onp.abs(max_range))
    quantized_range = 127.0
    scale = quantized_range / real_range
    assert qdata.dtype == onp.int8
    assert min_val.dtype == onp.float32
    assert max_val.dtype == onp.float32
    assert same(min_val.asscalar(), -real_range)
    assert same(max_val.asscalar(), real_range)
    qdata_np = (onp.sign(data_np) * onp.minimum(onp.abs(data_np) * scale + 0.5, quantized_range)).astype(onp.int8)
    assert_almost_equal(qdata.asnumpy(), qdata_np, atol = 1)

def test_calibrated_quantize_v2_bfloat16_to_int8():
    shape = rand_shape_nd(4)
    data = mx.nd.random.normal(0, 1, shape).astype('bfloat16')
    min_range = mx.nd.min(data).asscalar()
    max_range = mx.nd.max(data).asscalar()
    qdata, min_val, max_val = mx.nd.contrib.quantize_v2(data, 'int8', min_range, max_range)
    data_np = data.asnumpy()
    real_range = onp.maximum(onp.abs(min_range), onp.abs(max_range))
    quantized_range = 127.0
    scale = quantized_range / real_range
    assert qdata.dtype == onp.int8
    assert min_val.dtype == onp.float32
    assert max_val.dtype == onp.float32
    assert same(min_val.asscalar(), -real_range)
    assert same(max_val.asscalar(), real_range)
    qdata_np = (onp.sign(data_np) * onp.minimum(onp.abs(data_np) * scale + 0.5, quantized_range)).astype(onp.int8)
    assert_almost_equal(qdata.asnumpy(), qdata_np, atol=1)

def test_dequantize_int8_to_float32():

    def get_test_data(real_range, qdata_np):
        qdata = mx.nd.array(qdata_np, dtype=onp.int8)
        min_range = mx.nd.array([-real_range], dtype=onp.float32)
        max_range = mx.nd.array([real_range], dtype=onp.float32)
        return qdata, min_range, max_range

    def baseline_dequantization(qdata, real_range, qdata_np):
        quantized_range = 127.0
        scale = real_range / quantized_range
        data_np = qdata_np * scale
        return data_np

    def test_nd_array_dequantization(qdata, min_range, max_range, expected_result):
        data = mx.nd.contrib.dequantize(qdata, min_range, max_range, out_type='float32')
        assert data.dtype == onp.float32
        assert_almost_equal(data.asnumpy(), expected_result, atol = 1)

    def test_symbolic_api_dequantization(qdata, min_range, max_range, expected_result):
        sym_data = mx.sym.Variable('data')
        sym_min_range = mx.sym.Variable('min_range')
        sym_max_range = mx.sym.Variable('max_range')
        dequant = mx.sym.contrib.dequantize(sym_data, sym_min_range,
                                            sym_max_range, out_type='float32')
        out = dequant._bind(ctx=mx.current_device(),
                           args={'data':qdata, 'min_range':min_range, 'max_range':max_range})
        data = out.forward()[0]
        assert data.dtype == onp.float32
        assert_almost_equal(data.asnumpy(), expected_result, atol = 1)

    real_range = 128
    shape = rand_shape_nd(4)
    qdata_np = onp.random.uniform(low=-127, high=127, size=shape).astype(dtype=onp.int8)
    qdata, min_range, max_range = get_test_data(real_range, qdata_np)
    expected_result = baseline_dequantization(qdata, real_range, qdata_np)
    # test nd array implementation.
    test_nd_array_dequantization(qdata, min_range, max_range, expected_result)
    # test symbolic api implementaion.
    test_symbolic_api_dequantization(qdata, min_range, max_range, expected_result)


def test_requantize_int32_to_int8():
    def quantized_int32_to_float(qdata, min_range, max_range):
        assert qdata.dtype == 'int32'
        quantized_range = onp.iinfo('int32').max
        real_range = onp.maximum(onp.abs(min_range), onp.abs(max_range))
        scale = float(real_range) / float(quantized_range)
        return qdata.astype('float32') * scale

    def float_to_quantized_int8(data, min_range, max_range):
        assert data.dtype == 'float32'
        real_range = onp.maximum(onp.abs(min_range), onp.abs(max_range))
        quantized_range = onp.iinfo('int8').max
        scale = float(quantized_range) / float(real_range)
        return (onp.sign(data) * onp.minimum(onp.abs(data) * scale + 0.5, quantized_range)).astype('int8')

    def requantize(qdata, min_data, max_data, real_range):
        data = quantized_int32_to_float(qdata, min_data, max_data)
        output = float_to_quantized_int8(data, -real_range, real_range)
        return output, -real_range, real_range

    def requantize_baseline(qdata, min_data, max_data, min_calib_range=None, max_calib_range=None):
        if min_calib_range is not None and max_calib_range is not None:
            real_range = onp.maximum(onp.abs(min_calib_range), onp.abs(max_calib_range))
            return requantize(qdata, min_data, max_data, real_range)
        else:
            min_range = quantized_int32_to_float(onp.min(qdata), min_data, max_data)
            max_range = quantized_int32_to_float(onp.max(qdata), min_data, max_data)
            return requantize(qdata, min_data, max_data, onp.maximum(onp.abs(min_range), onp.abs(max_range)))

    def check_requantize(shape, min_calib_range=None, max_calib_range=None):
        qdata = mx.nd.random.uniform(low=-1000.0, high=1000.0, shape=shape).astype('int32')
        min_range = mx.nd.array([-1010.0])
        max_range = mx.nd.array([1020.0])
        if min_calib_range is None or max_calib_range is None:
            qdata_int8, min_output, max_output = mx.nd.contrib.requantize(qdata, min_range, max_range)
        else:
            qdata_int8, min_output, max_output = mx.nd.contrib.requantize(qdata, min_range, max_range,
                                                                          min_calib_range=min_calib_range,
                                                                          max_calib_range=max_calib_range)

        qdata_int8_np, min_output_np, max_output_np = requantize_baseline(qdata.asnumpy(), min_range.asscalar(),
                                                                          max_range.asscalar(),
                                                                          min_calib_range=min_calib_range,
                                                                          max_calib_range=max_calib_range)
        assert_almost_equal(qdata_int8.asnumpy(), qdata_int8_np, atol = 1)
        assert_almost_equal(min_output.asnumpy(), onp.array([min_output_np]))
        assert_almost_equal(max_output.asnumpy(), onp.array([max_output_np]))

    @use_np
    def check_requantize_with_gluon(shape, min_calib_range=None, max_calib_range=None):
        qdata = mx.np.random.uniform(low=-1000.0, high=1000.0, size=shape).astype('int32')
        min_range = mx.np.array([-1010.0])
        max_range = mx.np.array([1020.0])

        class RequantizeBlock(mx.gluon.nn.HybridBlock):
            def __init__(self, min_calib_range=None, max_calib_range=None, **kwargs):
                super(RequantizeBlock, self).__init__(**kwargs)
                self.min_calib_range = min_calib_range
                self.max_calib_range = max_calib_range

            def forward(self, x, min_range, max_range):
                if self.min_calib_range is not None and self.max_calib_range is not None:
                    out = npx.requantize(x, min_range, max_range,
                                         min_calib_range=self.min_calib_range,
                                         max_calib_range=self.max_calib_range)
                else:
                    out = npx.requantize(x, min_range, max_range)
                return out

        requant = RequantizeBlock(min_calib_range, max_calib_range)  # m*_calib_ranges can be None
        qdata_int8, min_output, max_output = requant(qdata, min_range, max_range)

        qdata_int8_np, min_output_np, max_output_np = requantize_baseline(qdata.asnumpy(), min_range.item(),
                                                                          max_range.item(),
                                                                          min_calib_range=min_calib_range,
                                                                          max_calib_range=max_calib_range)
        assert_almost_equal(qdata_int8.asnumpy(), qdata_int8_np, atol = 1)
        assert_almost_equal(min_output.asnumpy(), onp.array([min_output_np]))
        assert_almost_equal(max_output.asnumpy(), onp.array([max_output_np]))

    # test with gluon API.
    check_requantize_with_gluon((3, 4, 10, 10))
    check_requantize_with_gluon((32, 3, 23, 23))
    check_requantize_with_gluon((3, 4, 10, 10), min_calib_range=-1050.0, max_calib_range=1040.0)
    check_requantize_with_gluon((32, 3, 23, 23), min_calib_range=-134.349, max_calib_range=523.43)
    # Test with nd array API
    check_requantize((3, 4, 10, 10))
    check_requantize((32, 3, 23, 23))
    check_requantize((3, 4, 10, 10), min_calib_range=-1050.0, max_calib_range=1040.0)
    check_requantize((32, 3, 23, 23), min_calib_range=-134.349, max_calib_range=523.43)


@use_np
def test_quantized_conv():
    def check_quantized_conv(data_shape, kernel, num_filter, pad, stride, dilate, use_bias, qdtype):
        if is_test_for_native_cpu():
            print('skipped testing quantized_conv for native cpu since it is not supported yet')
            return
        elif is_test_for_dnnl():
            # (TODO)Xinyu: https://github.com/apache/mxnet/issues/16830
            print('skipped testing quantized_conv for oneDNN cpu since it is a flaky case')
            return
        elif qdtype == 'uint8' and is_test_for_gpu():
            print('skipped testing quantized_conv for gpu uint8 since it is not supported yet')
            return
        elif is_test_for_gpu() and len(data_shape) != 4:
            print('skipped testing quantized_conv for gpu 5d layout since it is not supported yet')
            return

        # run fp32 conv
        if len(data_shape) == 4:
            convfp32 = mx.gluon.nn.Conv2D(channels=num_filter, kernel_size=kernel, strides=stride,
                                          padding=pad, dilation=dilate, use_bias=use_bias)
        elif len(data_shape) == 5:
            convfp32 = mx.gluon.nn.Conv3D(channels=num_filter, kernel_size=kernel, strides=stride,
                                          padding=pad, dilation=dilate, use_bias=use_bias)
        else:
            print('unsupported shape')
            assert False

        if qdtype == 'uint8':
            data_low = 0.0
            data_high = 127.0
        else:
            data_low = -127.0
            data_high = 127.0

        convfp32.initialize()
        input_data = mx.np.random.uniform(low=data_low,
                                          high=data_high,
                                          size=data_shape
                                         ).astype('int32').astype('float32')
        convfp32(input_data) # initialize params
        npx.waitall()
        fp32_params = convfp32.collect_params()
        weight_shape = fp32_params['weight'].shape
        new_args = dict()
        new_args['weight'] = mx.np.random.uniform(low=-127.0,
                                                  high=127.0,
                                                  size=weight_shape
                                                 ).astype('int32').astype('float32')
        if use_bias:
           new_args['bias'] = mx.np.random.uniform(low=-127.0,
                                                   high=127.0,
                                                   size=fp32_params['bias'].shape
                                                  ).astype('int32').astype('float32')
        convfp32.load_dict(new_args, cast_dtype=True, dtype_source='saved')

        output = convfp32(input_data)

        # run quantized conv
        class QuantConv(mx.gluon.nn.HybridBlock):
            def __init__(self, channels, kernel_size, strides=(1, 1),
                         padding=(0, 0), dilation=(1, 1), use_bias=True, **kwargs):
                super(QuantConv, self).__init__(**kwargs)
                self.use_bias = use_bias
                self._kwargs = {'kernel': kernel_size, 'stride': strides, 'dilate': dilation,
                                'pad': padding, 'num_filter': channels, 'no_bias': not use_bias, 'num_group': 1,
                                'layout': 'NCHW'}

                self.min_data = mx.gluon.Parameter('min_data', dtype='float32', shape=(1), allow_deferred_init=True)
                self.max_data = mx.gluon.Parameter('max_data', dtype='float32', shape=(1), allow_deferred_init=True)

                self.weight = mx.gluon.Parameter('weight', dtype='int8', shape=weight_shape, allow_deferred_init=True)
                self.min_weight = mx.gluon.Parameter('min_weight', dtype='float32', shape=(1), allow_deferred_init=True)
                self.max_weight = mx.gluon.Parameter('max_weight', dtype='float32', shape=(1), allow_deferred_init=True)

                if use_bias:
                    self.bias = mx.gluon.Parameter('bias', dtype='int8', shape=(num_filter,), allow_deferred_init=True)
                    self.min_bias = mx.gluon.Parameter('min_bias', dtype='float32', shape=(1), allow_deferred_init=True)
                    self.max_bias = mx.gluon.Parameter('max_bias', dtype='float32', shape=(1), allow_deferred_init=True)

            def forward(self, x):
                device = x.device
                weight = self.weight.data().to_device(device)
                bias = self.bias.data().to_device(device) if self.use_bias else None
                min_data = self.min_data.data().to_device(device)
                max_data = self.max_data.data().to_device(device)
                min_weight = self.min_weight.data().to_device(device)
                max_weight = self.max_weight.data().to_device(device)
                min_bias = self.min_bias.data().to_device(device) if self.use_bias else None
                max_bias = self.max_bias.data().to_device(device) if self.use_bias else None
                out = npx.quantized_conv(data=x, weight=weight, bias=bias,
                                         min_data=min_data, max_data=max_data,
                                         min_weight=min_weight, max_weight=max_weight,
                                         min_bias=min_bias, max_bias=max_bias,
                                         **self._kwargs)
                return out

        convint8 = QuantConv(channels=num_filter, kernel_size=kernel, strides=stride,
                             padding=pad, dilation=dilate, use_bias=use_bias)

        quantized_range = 127.0
        qargs = {
            'weight': new_args['weight'].astype('int8'),
            'min_data': mx.np.array([-quantized_range]),
            'max_data': mx.np.array([quantized_range]),
            'min_weight': mx.np.array([-quantized_range]),
            'max_weight': mx.np.array([quantized_range])
        }
        if use_bias:
            qargs.update({
                'bias': new_args['bias'].astype('int8'),
                'min_bias': mx.np.array([-quantized_range]),
                'max_bias': mx.np.array([quantized_range]),
            })

        convint8.load_dict(qargs, cast_dtype=True, dtype_source='saved')

        qoutput, min_range, max_range = convint8(input_data.astype(qdtype))

        if use_bias:
            # with adding bias, accuracy loss should not be greater than one
            diff = mx.np.abs(output - qoutput.astype(output.dtype))
            cond = mx.np.less(2, diff).sum().item()
            assert cond == 0
        else:
            assert_almost_equal(output.asnumpy(), qoutput.asnumpy(), atol = 1)

    for qdtype in ['int8', 'uint8']:
        check_quantized_conv((3, 4, 28, 28), (3, 3), 128, (1, 1), (1, 1), (1, 1), True, qdtype)
        check_quantized_conv((3, 4, 28, 28), (3, 3), 128, (1, 1), (1, 1), (1, 1), False, qdtype)
        check_quantized_conv((1, 3, 4, 28, 28), (1, 3, 3), 128, (1, 1, 1), (1, 1, 1), (1, 1, 1), False, qdtype)
        check_quantized_conv((1, 3, 4, 28, 28), (1, 3, 3), 128, (1, 1, 1), (1, 1, 1), (1, 1, 1), True, qdtype)
        check_quantized_conv((1, 3, 4, 28, 28), (1, 3, 3), 128, (1, 1, 1), (1, 1, 1), (2, 2, 2), False, qdtype)
        check_quantized_conv((1, 3, 4, 28, 28), (1, 3, 3), 128, (1, 1, 1), (1, 1, 1), (2, 2, 2), True, qdtype)


@use_np
def test_quantized_elemwise_add():
    def check_quantized_elemwise_add(data_shape, qdtypeA, qdtypeB):
        if is_test_for_native_cpu():
            print('skipped testing quantized_elemwise_add for native cpu since it is not supported yet')
            return
        elif (qdtypeA != 'uint8' and qdtypeA != 'int8') or (qdtypeB != 'uint8' and qdtypeB != 'int8'):
            print('skipped testing quantized_elemwise_add for not supported data type')
            return
        elif is_test_for_gpu():
            print('skipped testing quantized_elemwise_add for gpu since it is not supported yet')
            return

        class ElemwiseSumBlock(mx.gluon.nn.HybridBlock):
            def __init__(self, **kwargs):
                super(ElemwiseSumBlock, self).__init__(**kwargs)

            def forward(self, dataA, dataB):
                return dataA + dataB

        class QuantElemwiseSumBlock(mx.gluon.nn.HybridBlock):
            def __init__(self, **kwargs):
                super(QuantElemwiseSumBlock, self).__init__(**kwargs)

            def forward(self, dataA, dataB, dataA_min, dataA_max, dataB_min, dataB_max):
                return npx.quantized_elemwise_add(dataA, dataB, dataA_min, dataA_max, dataB_min, dataB_max)

        elemwise_add_fp32 = ElemwiseSumBlock()

        dataA_low, dataA_high = get_low_high(qdtypeA)
        dataB_low, dataB_high = get_low_high(qdtypeB)

        dataA_val = mx.np.random.uniform(low=dataA_low, high=dataA_high, size=data_shape).astype('int32').astype('float32')
        dataB_val = mx.np.random.uniform(low=dataB_low, high=dataB_high, size=data_shape).astype('int32').astype('float32')

        output = elemwise_add_fp32(dataA_val, dataB_val)
        mx.nd.waitall()

        #run quantized
        quantized_elemwise_add = QuantElemwiseSumBlock()
        dataA_val_int8 = dataA_val.astype(qdtypeA)
        dataB_val_int8 = dataB_val.astype(qdtypeB)

        quantized_range = 127.0
        min_dataA = mx.np.array([dataA_low])
        max_dataA = mx.np.array([dataA_high])
        min_dataB = mx.np.array([dataB_low])
        max_dataB = mx.np.array([dataB_high])
        qoutput, min_range, max_range = quantized_elemwise_add(dataA_val_int8, dataB_val_int8,
                                                               min_dataA, max_dataA,
                                                               min_dataB, max_dataB)
        int8_rslt = qoutput.astype(output.dtype) * max_range / 0x7fffffff
        diff = mx.np.abs(output - int8_rslt)
        cond = mx.np.less(2, diff).sum().item()
        assert cond == 0

    check_quantized_elemwise_add((4, 6), 'uint8', 'int8')
    check_quantized_elemwise_add((13, 74, 52), 'uint8', 'uint8')
    check_quantized_elemwise_add((3, 4, 56, 56), 'int8', 'uint8')
    check_quantized_elemwise_add((32, 56, 64, 11), 'int8', 'int8')

@use_np
def test_quantized_npi_add():
    def check_quantized_npi_add(data_shape,  qdtypeA, qdtypeB, broadcast=None):
        if is_test_for_native_cpu():
            print('skipped testing quantized_npi_add for native cpu since it is not supported yet')
            return
        elif (qdtypeA != 'uint8' and qdtypeA != 'int8') or (qdtypeB != 'uint8' and qdtypeB != 'int8'):
            print('skipped testing quantized_npi_add for not supported data type')
            return
        elif is_test_for_gpu():
            print('skipped testing quantized_npi_add for gpu since it is not supported yet')
            return

        class ElemwiseSumBlock(mx.gluon.nn.HybridBlock):
            def __init__(self, **kwargs):
                super(ElemwiseSumBlock, self).__init__(**kwargs)

            def forward(self, dataA, dataB):
                return dataA + dataB

        class QuantElemwiseSumBlock(mx.gluon.nn.HybridBlock):
            def __init__(self, **kwargs):
                super(QuantElemwiseSumBlock, self).__init__(**kwargs)

            def forward(self, dataA, dataB, dataA_min, dataA_max, dataB_min, dataB_max):
                return npx.quantized_npi_add(dataA, dataB, dataA_min, dataA_max, dataB_min, dataB_max)

        elemwise_add_fp32 = ElemwiseSumBlock()

        dataA_low, dataA_high = get_low_high(qdtypeA)
        dataB_low, dataB_high = get_low_high(qdtypeB)

        data_shapeA = data_shape
        data_shapeB = data_shape

        if broadcast :
            if broadcast == 'A':
                data_shapeA = ()
                for index in range(len(data_shape)):
                    data_shapeA += (1,)
            else:
                data_shapeB = ()
                for index in range(len(data_shape)):
                    data_shapeB += (1,)

        dataA_val = mx.np.random.uniform(low=dataA_low, high=dataA_high, size=data_shapeA).astype('int32').astype('float32')
        dataB_val = mx.np.random.uniform(low=dataB_low, high=dataB_high, size=data_shapeB).astype('int32').astype('float32')

        output = elemwise_add_fp32(dataA_val, dataB_val)

        #run quantized
        quantized_elemwise_add = QuantElemwiseSumBlock()
        dataA_val_int8 = dataA_val.astype(qdtypeA)
        dataB_val_int8 = dataB_val.astype(qdtypeB)
        quantized_range = 127.0
        min_dataA = mx.np.array([dataA_low])
        max_dataA = mx.np.array([dataA_high])
        min_dataB = mx.np.array([dataB_low])
        max_dataB = mx.np.array([dataB_high])
        qoutput, min_range, max_range = quantized_elemwise_add(dataA_val_int8, dataB_val_int8,
                                                               min_dataA, max_dataA,
                                                               min_dataB, max_dataB)
        int8_rslt = qoutput.astype(output.dtype) * max_range / 0x7fffffff
        diff = mx.np.abs(output - int8_rslt)
        cond = mx.np.less(2, diff).sum().item()
        assert cond == 0

    check_quantized_npi_add((4, 6), 'uint8', 'int8')
    check_quantized_npi_add((13, 74, 52), 'uint8', 'uint8')
    check_quantized_npi_add((3, 4, 56, 56), 'int8', 'uint8')
    check_quantized_npi_add((32, 56, 64, 11), 'int8', 'int8')

    check_quantized_npi_add((4, 6), 'uint8', 'int8', 'A')
    check_quantized_npi_add((13, 74, 52), 'uint8', 'uint8', 'B')
    check_quantized_npi_add((3, 4, 56, 56), 'int8', 'uint8', 'A')
    check_quantized_npi_add((32, 56, 64, 11), 'int8', 'int8', 'B')


@use_np
def test_quantized_elemwise_mul():
    def check_quantized_elemwise_mul(data_shape, qtype):
        if is_test_for_native_cpu():
            print('skipped testing quantized_elemwise_mul for native cpu since it is not supported yet')
            return
        elif qtype != 'int8':
            print('skipped testing quantized_elemwise_mul for not supported data type')
            return
        elif is_test_for_gpu():
            print('skipped testing quantized_elemwise_mul for gpu since it is not supported yet')
            return

        class ElemwiseMulBlock(mx.gluon.nn.HybridBlock):
            def __init__(self, **kwargs):
                super(ElemwiseMulBlock, self).__init__(**kwargs)

            def forward(self, dataA, dataB):
                return mx.np.multiply(dataA, dataB)

        class QuantElemwiseMulBlock(mx.gluon.nn.HybridBlock):
            def __init__(self, **kwargs):
                super(QuantElemwiseMulBlock, self).__init__(**kwargs)

            def forward(self, dataA, dataB, dataA_min, dataA_max, dataB_min, dataB_max):
                return npx.quantized_elemwise_mul(dataA, dataB, dataA_min, dataA_max, dataB_min, dataB_max)

        elemwise_mul_fp32 = ElemwiseMulBlock()
        data_low, data_high = get_low_high(qtype)

        dataA_val = mx.np.random.uniform(low=data_low, high=data_high, size=data_shape).astype('int32').astype('float32')
        dataB_val = mx.np.random.uniform(low=data_low, high=data_high, size=data_shape).astype('int32').astype('float32')

        output = elemwise_mul_fp32(dataA_val, dataB_val)

        quantized_elemwise_mul = QuantElemwiseMulBlock()
        dataA_val_int8 = dataA_val.astype(qtype)
        dataB_val_int8 = dataB_val.astype(qtype)
        quantized_range = 127.0
        min_dataA = mx.np.array([data_low])
        max_dataA = mx.np.array([data_high])
        min_dataB = mx.np.array([data_low])
        max_dataB = mx.np.array([data_high])
        qoutput, min_range, max_range = quantized_elemwise_mul(dataA_val_int8, dataB_val_int8,
                                                               min_dataA, max_dataA,
                                                               min_dataB, max_dataB)

        fp32_rslt = output.asnumpy()
        int8_rslt = qoutput.astype(output.dtype)
        assert_almost_equal(fp32_rslt, int8_rslt, atol = 1e-4)

    for qtype in ['int8', 'uint8']:
        check_quantized_elemwise_mul((4, 6), qtype)
        check_quantized_elemwise_mul((13, 74, 52), qtype)
        check_quantized_elemwise_mul((3, 4, 56, 56), qtype)
        check_quantized_elemwise_mul((32, 56, 64, 11), qtype)


@use_np
def test_quantized_pooling():
    def check_quantized_pooling(data_shape, kernel, pool_type, pad, stride, global_pool, qdtype, convention='valid'):
        if is_test_for_native_cpu():
            print('skipped testing quantized_pooling for native cpu since it is not supported yet')
            return
        elif qdtype == 'uint8' and is_test_for_gpu():
            print('skipped testing quantized_pooling for gpu uint8 since it is not supported yet')
            return
        elif is_test_for_gpu() and len(data_shape) != 4:
            print('skipped testing quantized_pooling for gpu 5d layout since it is not supported yet')
            return

        class PoolingBlock(mx.gluon.nn.HybridBlock):
            def __init__(self, kernel=kernel, pad=pad, stride=stride,
                         pool_type=pool_type, global_pool=global_pool, cudnn_off=False,
                         pooling_convention=convention):
                super(PoolingBlock, self).__init__()
                self._kwargs = {'kernel': kernel, 'pad': pad, 'stride': stride,
                                'pool_type': pool_type, 'global_pool': global_pool,
                                'cudnn_off': False, 'pooling_convention': convention}

            def forward(self, data):
                return npx.pooling(data, **self._kwargs)

        class QuantPoolingBlock(mx.gluon.nn.HybridBlock):
            def __init__(self, kernel=kernel, pad=pad, stride=stride,
                         pool_type=pool_type, global_pool=global_pool,
                         cudnn_off=False, pooling_convention=convention):
                super(QuantPoolingBlock, self).__init__()

                self._kwargs = {'kernel': kernel, 'pad': pad, 'stride': stride,
                                'pool_type': pool_type, 'global_pool': global_pool, 'cudnn_off': False,
                                'pooling_convention':convention}

            def forward(self, data, min_data, max_data):
                return npx.quantized_pooling(data, min_data, max_data, **self._kwargs)

        pooling_fp32 = PoolingBlock()
        if qdtype == 'uint8':
            data_low = 0.0
            data_high = 127.0
        else:
            data_low = -127.0
            data_high = 127.0

        input_data = mx.np.random.uniform(low=data_low,
                                          high=data_high,
                                          size=data_shape
                                         ).astype('int32').astype('float32')
        output = pooling_fp32(input_data)

        quantized_pooling = QuantPoolingBlock(kernel=kernel, pad=pad, stride=stride,
                                              pool_type=pool_type, global_pool=global_pool,
                                              pooling_convention=convention)

        int8_input_data = input_data.astype(qdtype)
        quantized_range = 127.0
        min_data = mx.np.array([-quantized_range])
        max_data = mx.np.array([quantized_range])

        qoutput, min_range, max_range = quantized_pooling(int8_input_data, min_data, max_data)

        if pool_type == 'max':
            assert_almost_equal(output.asnumpy(), qoutput.asnumpy())
        elif pool_type == 'avg':  # for avg pooling, fp32 and int8 may be different due to rounding errors
            diff = mx.np.abs(output - qoutput.astype(output.dtype))
            cond = mx.np.less(2, diff).sum().item()
            assert cond == 0

    for qdtype in ['int8', 'uint8']:
        check_quantized_pooling((3, 4, 56, 56), (3, 3), 'max', (0, 0), (2, 2), False, qdtype)
        check_quantized_pooling((3, 4, 56, 56), (3, 3), 'max', (0, 0), (2, 2), True, qdtype)
        check_quantized_pooling((3, 512, 7, 7), (7, 7), 'avg', (0, 0), (1, 1), False, qdtype)
        check_quantized_pooling((3, 512, 7, 7), (7, 7), 'avg', (0, 0), (1, 1), True, qdtype)
        check_quantized_pooling((3, 4, 3, 56, 56), (1, 3, 3), 'max', (0, 0, 0), (1, 2, 2), False, qdtype)
        check_quantized_pooling((3, 4, 3, 56, 56), (1, 3, 3), 'max', (0, 0, 0), (1, 2, 2), True, qdtype)
        check_quantized_pooling((3, 512, 3, 7, 7), (1, 7, 7), 'avg', (0, 0, 0), (1, 2, 2), False, qdtype)
        check_quantized_pooling((3, 512, 3, 7, 7), (1, 7, 7), 'avg', (0, 0, 0), (1, 2, 2), True, qdtype)

        check_quantized_pooling((3, 4, 56, 56), (3, 3), 'max', (0, 0), (2, 2), False, qdtype, 'full')
        check_quantized_pooling((3, 4, 56, 56), (3, 3), 'max', (0, 0), (2, 2), True, qdtype, 'full')
        check_quantized_pooling((3, 512, 7, 7), (7, 7), 'avg', (0, 0), (1, 1), False, qdtype, 'full')
        check_quantized_pooling((3, 512, 7, 7), (7, 7), 'avg', (0, 0), (1, 1), True, qdtype, 'full')
        check_quantized_pooling((3, 4, 3, 56, 56), (1, 3, 3), 'max', (0, 0, 0), (1, 2, 2), False, qdtype, 'full')
        check_quantized_pooling((3, 4, 3, 56, 56), (1, 3, 3), 'max', (0, 0, 0), (1, 2, 2), True, qdtype, 'full')
        check_quantized_pooling((3, 512, 3, 7, 7), (1, 7, 7), 'avg', (0, 0, 0), (1, 2, 2), False, qdtype, 'full')
        check_quantized_pooling((3, 512, 3, 7, 7), (1, 7, 7), 'avg', (0, 0, 0), (1, 2, 2), True, qdtype, 'full')


@use_np
def test_quantized_fc():
    def check_quantized_fc(data_shape, num_hidden, use_bias, qdtype, flatten=True):
        if is_test_for_native_cpu():
            hasMKL = False
            for key in os.environ.keys():
                if operator.eq(key, "BUILD_TAG"):
                    if os.environ['BUILD_TAG'].find("MKL") != -1:
                        hasMKL = True
                    break
            if hasMKL == False:
                print('skipped testing quantized_fc on cpu since s8u8s32 is only supported by MKL BLAS library')
                return
        elif qdtype == 'uint8' and is_test_for_gpu():
            print('skipped testing quantized_fc for gpu uint8 since it is not supported yet')
            return

        def maxabs(a, b):
            return mx.np.maximum(mx.np.abs(a), mx.np.abs(b))

        int8_range = 127.0
        if qdtype == 'uint8':
            data_low = 0.0
            data_high = 63.0
            quantized_range = 255.0
        else:
            data_low = -63.0
            data_high = 63.0
            quantized_range = 127.0

        data = mx.np.random.uniform(low=data_low,
                                    high=data_high,
                                    size=data_shape
                                   ).astype('int32').astype('float32')
        fc_fp32 = mx.gluon.nn.Dense(units=num_hidden, use_bias=use_bias, flatten=flatten)
        fc_fp32.initialize()
        fc_fp32(data)
        npx.waitall()
        fp32_params = fc_fp32.collect_params()
        weight_shape = fp32_params['weight'].shape

        new_args = dict()
        new_args['weight'] = mx.np.random.uniform(low=data_low,
                                                  high=data_high,
                                                  size=fp32_params['weight'].shape
                                                 ).astype('int32').astype('float32')
        data_min = mx.np.min(data).astype('float32')
        data_max = mx.np.max(data).astype('float32')
        weight_min = mx.np.min(new_args['weight']).astype('float32')
        weight_max = mx.np.max(new_args['weight']).astype('float32')
        data_range = maxabs(data_min, data_max)
        weight_range = maxabs(weight_min, weight_max)

        if use_bias:
            bias = mx.np.random.uniform(low=data_low,
                                        high=data_high,
                                        size=fp32_params['bias'].shape
                                       ).astype('int32').astype('float32')
            bias_min = mx.np.min(bias).astype('float32')
            bias_max = mx.np.max(bias).astype('float32')
            bias_range = maxabs(bias_min, bias_max)

            bias_scale = int8_range / bias_range
            data_scale = quantized_range / data_range
            weight_scale = int8_range / weight_range
            bias_int32_rescale = data_scale * weight_scale / bias_scale
            new_bias = bias.astype('float32') * bias_int32_rescale
            new_args['bias'] = new_bias.astype('int32').astype('float32')

        fc_fp32.load_dict(new_args, cast_dtype=True, dtype_source='saved')
        output = fc_fp32(data)

        class QuantFC(mx.gluon.nn.HybridBlock):
            def __init__(self, num_hidden, use_bias, flatten, **kwargs):
                super(QuantFC, self).__init__(**kwargs)
                self.use_bias = use_bias
                self._kwargs = {'num_hidden': num_hidden, 'no_bias': not use_bias, 'flatten': flatten}

                self.min_data = mx.gluon.Parameter('min_data', dtype='float32', shape=(1), allow_deferred_init=True)
                self.max_data = mx.gluon.Parameter('max_data', dtype='float32', shape=(1), allow_deferred_init=True)

                self.weight = mx.gluon.Parameter('weight', dtype='int8', shape=weight_shape, allow_deferred_init=True)
                self.min_weight = mx.gluon.Parameter('min_weight', dtype='float32', shape=(1), allow_deferred_init=True)
                self.max_weight = mx.gluon.Parameter('max_weight', dtype='float32', shape=(1), allow_deferred_init=True)

                if use_bias:
                    self.bias = mx.gluon.Parameter('bias', dtype='int8', shape=(num_hidden,), allow_deferred_init=True)
                    self.min_bias = mx.gluon.Parameter('min_bias', dtype='float32', shape=(1), allow_deferred_init=True)
                    self.max_bias = mx.gluon.Parameter('max_bias', dtype='float32', shape=(1), allow_deferred_init=True)

            def forward(self, x):
                device = x.device
                weight = self.weight.data().to_device(device)
                bias = self.bias.data().to_device(device) if self.use_bias else None
                min_data = self.min_data.data().to_device(device)
                max_data = self.max_data.data().to_device(device)
                min_weight = self.min_weight.data().to_device(device)
                max_weight = self.max_weight.data().to_device(device)
                min_bias = self.min_bias.data().to_device(device) if self.use_bias else None
                max_bias = self.max_bias.data().to_device(device) if self.use_bias else None
                out = npx.quantized_fully_connected(data=x, weight=weight, bias=bias,
                                                    min_data=min_data, max_data=max_data,
                                                    min_weight=min_weight, max_weight=max_weight,
                                                    min_bias=min_bias, max_bias=max_bias,
                                                    **self._kwargs)
                return out

        fc_int8 = QuantFC(num_hidden=num_hidden, use_bias=use_bias, flatten=flatten)
        qargs = {
            'weight': new_args['weight'].astype('int8'),
            'min_data': mx.np.array([-data_range]),
            'max_data': mx.np.array([data_range]),
            'min_weight': mx.np.array([-weight_range]),
            'max_weight': mx.np.array([weight_range])
        }
        if use_bias:
            qargs.update({
                'bias': bias.astype('int8'),
                'min_bias': mx.np.array([-bias_range]),
                'max_bias': mx.np.array([bias_range]),
            })

        fc_int8.load_dict(qargs, cast_dtype=True, dtype_source='saved')

        qoutput, min_range, max_range = fc_int8(data.astype(qdtype))

        if use_bias:
            # with adding bias, accuracy loss should not be greater than one
            diff = mx.np.abs(output - qoutput.astype(output.dtype))
            cond = mx.np.less(2, diff).sum().item()
            assert cond == 0
        else:
            assert_almost_equal(output.asnumpy(), qoutput.asnumpy())

    for qdtype in ['int8', 'uint8']:
        if is_test_for_dnnl():
            check_quantized_fc((32, 512, 2), 100, False, qdtype, flatten=False)
            check_quantized_fc((32, 512, 2), 100, True, qdtype, flatten=False)
            check_quantized_fc((32, 512, 2, 2), 100, False, qdtype, flatten=False)
            check_quantized_fc((32, 512, 2, 2), 100, True, qdtype, flatten=False)
        check_quantized_fc((32, 512, 2, 2), 100, False, qdtype)
        check_quantized_fc((32, 111, 2, 2), 100, False, qdtype)
        check_quantized_fc((32, 512, 2, 2), 100, True, qdtype)
        check_quantized_fc((32, 111, 2, 2), 100, True, qdtype)
        check_quantized_fc((256, 2048, 2, 2), 800, True, qdtype)
        check_quantized_fc((256, 111, 2, 2), 800, True, qdtype)
        check_quantized_fc((256, 2048, 2, 2), 800, False, qdtype)
        check_quantized_fc((256, 111, 2, 2), 800, False, qdtype)

@use_np
def test_quantized_transpose():
    def check_quantized_transpose(shape, qdtype, axes):
        data_low, data_high = get_low_high(qdtype)
        data = mx.np.random.uniform(low=data_low, high=data_high, size=shape).astype(qdtype).astype('float32')
        min_data = mx.np.array([mx.np.min(data).astype('float32').item()])
        max_data = mx.np.array([mx.np.max(data).astype('float32').item()])
        qdata = data.astype(qdtype)
        output = mx.np.transpose(data, axes=axes)
        qoutput, min_output, max_output = npx.quantized_transpose(qdata, min_data, max_data, axes=axes)
        assert_almost_equal(output.asnumpy(), qoutput.asnumpy())
        assert_almost_equal(min_output.item(), min_data.item())
        assert_almost_equal(max_output.item(), max_data.item())

    for qtype in ['int8', 'uint8']:
        check_quantized_transpose((), qtype, ())
        check_quantized_transpose((2,3), qtype, (1,0))
        check_quantized_transpose((8,21), qtype, (1,0))
        check_quantized_transpose((7,3,9), qtype, (2,1,0))
        check_quantized_transpose((5,3,6,8), qtype, (2,3,0,1))


@use_np
def test_quantized_embedding():
    def check_quantized_embedding(data_shape, input_dim, output_dim):
        if is_test_for_gpu():
            print('skipped testing test_quantized_embedding for gpu since it is not supported yet')
            return

        def maxabs(a, b):
            return mx.np.maximum(mx.np.abs(a), mx.np.abs(b))

        data = mx.np.random.uniform(low=0,
                                    high=input_dim,
                                    size=data_shape
                                   ).astype('int32').astype('float32')
        embedding_fp32 = mx.gluon.nn.Embedding(input_dim=input_dim, output_dim=output_dim)
        embedding_fp32.initialize()
        embedding_fp32(data)
        npx.waitall()
        fp32_params = embedding_fp32.collect_params()
        weight_shape = fp32_params['weight'].shape
        int8_range = 127.0
        new_params = dict()
        weight = mx.np.random.uniform(low=-int8_range,
                                      high=int8_range,
                                      size=weight_shape
                                     ).astype('int32').astype('float32')
        new_params['weight'] = weight
        embedding_fp32.load_dict(new_params, cast_dtype=True, dtype_source='saved')

        output = embedding_fp32(data)

        weight_min = mx.np.min(weight).astype('float32')
        weight_max = mx.np.max(weight).astype('float32')
        weight_range = maxabs(weight_min, weight_max)

        class QuantEmbedding(mx.gluon.nn.HybridBlock):
            def __init__(self, input_dim=input_dim, output_dim=output_dim, **kwargs):
                super(QuantEmbedding, self).__init__(**kwargs)
                self._kwargs = {'input_dim': input_dim, 'output_dim': output_dim}

                self.weight = mx.gluon.Parameter('weight', dtype='float32', shape=weight_shape, allow_deferred_init=True)
                self.min_weight = mx.gluon.Parameter('min_weight', dtype='float32', shape=(1), allow_deferred_init=True)
                self.max_weight = mx.gluon.Parameter('max_weight', dtype='float32', shape=(1), allow_deferred_init=True)

            def forward(self, x):
                device = x.device
                weight = self.weight.data().to_device(device)
                min_weight = self.min_weight.data().to_device(device)
                max_weight = self.max_weight.data().to_device(device)
                out = npx.quantized_embedding(data=x, weight=weight,
                                              min_weight=min_weight,
                                              max_weight=max_weight,
                                              **self._kwargs)
                return out

        embedding_int8 = QuantEmbedding(input_dim=input_dim, output_dim=output_dim)
        qargs = {
            'weight': weight.astype('int8'),
            'min_weight': mx.np.array([-weight_range]),
            'max_weight': mx.np.array([weight_range])
        }

        embedding_int8.load_dict(qargs, cast_dtype=True, dtype_source='saved')

        qoutput, min_range, max_range = embedding_int8(data)


        assert_almost_equal(output.asnumpy(), qoutput.asnumpy())

    check_quantized_embedding((1,), 1000, 256)
    check_quantized_embedding((1,), 1024, 512)
    check_quantized_embedding((32,), 1000, 256)
    check_quantized_embedding((32,), 1024, 512)


@use_np
def test_quantized_flatten():
    def check_quantized_flatten(shape, qdtype):
        if qdtype == 'uint8':
            data_low = 0.0
            data_high = 127.0
        else:
            data_low = -127.0
            data_high = 127.0
        qdata = mx.np.random.uniform(low=data_low, high=data_high, size=shape).astype(qdtype)
        min_data = mx.np.array([-1023.343], dtype='float32')
        max_data = mx.np.array([2343.324275], dtype='float32')
        qoutput, min_output, max_output = npx.quantized_flatten(qdata, min_data, max_data)
        assert qoutput.ndim == 2
        assert qoutput.shape[0] == qdata.shape[0]
        assert qoutput.shape[1] == onp.prod(qdata.shape[1:])
        assert same(qdata.asnumpy().flatten(), qoutput.asnumpy().flatten())
        assert same(min_data.asnumpy(), min_output.asnumpy())
        assert same(max_data.asnumpy(), max_output.asnumpy())

    for qdtype in ['int8', 'uint8']:
        check_quantized_flatten((10,), qdtype)
        check_quantized_flatten((10, 15), qdtype)
        check_quantized_flatten((10, 15, 18), qdtype)
        check_quantized_flatten((3, 4, 23, 23), qdtype)


@use_np
def test_quantized_act():
    def check_quantized_act(data_shape, qdtype):
        if is_test_for_native_cpu():
            print('skipped testing quantized_act for native cpu since it is not supported yet')
            return
        elif qdtype == 'int8' and is_test_for_dnnl():
            print('skipped testing quantized_act for oneDNN cpu int8 since it is not supported yet')
            return
        elif is_test_for_gpu():
            print('skipped testing quantized_act for gpu since it is not supported yet')
            return

        act_fp32 = mx.gluon.nn.Activation(activation='relu')

        if qdtype == 'uint8':
            data_low = 0.0
            data_high = 127.0
        else:
            data_low = -127.0
            data_high = 127.0

        data = mx.np.random.uniform(low=data_low,
                                    high=data_high,
                                    size=data_shape
                                   ).astype(qdtype).astype('float32')
        output = act_fp32(data)

        class QuantActivation(mx.gluon.nn.HybridBlock):
            def __init__(self, activation, **kwargs):
                super(QuantActivation, self).__init__(**kwargs)
                self._kwargs = {'act_type': activation}

            def forward(self, x, min_data, max_data):
                out = npx.quantized_act(data=x, min_data=min_data, max_data=max_data, **self._kwargs)
                return out

        quantized_act = QuantActivation(activation='relu')

        qdata = data.astype(qdtype)
        quantized_range_min = mx.np.array([mx.np.min(data).astype('float32').item()])
        quantized_range_max = mx.np.array([mx.np.max(data).astype('float32').item()])
        qoutput, min_range, max_range = quantized_act(qdata, quantized_range_min, quantized_range_max)

        assert_almost_equal(output.asnumpy(), qoutput.asnumpy())
        assert_almost_equal(min_range.item(), quantized_range_min.item())
        assert_almost_equal(max_range.item(), quantized_range_max.item())

    for qdtype in ['int8', 'uint8']:
        check_quantized_act((10,), qdtype)
        check_quantized_act((10, 15), qdtype)
        check_quantized_act((10, 15, 18), qdtype)
        check_quantized_act((3, 4, 23, 23), qdtype)


@use_np
def test_quantized_bn():
    def get_mean_var(data):
        axes = list(range(data.ndim))
        del axes[1]
        mean = mx.np.mean(data, axis=axes)
        mean_broad = mx.np.expand_dims(mean, axis=0)
        mean_broad = mx.np.expand_dims(mean_broad, axis=2)
        mean_broad = mx.np.expand_dims(mean_broad, axis=3)
        mean_broad = mx.npx.broadcast_like(mean_broad, data)
        var = mx.np.multiply(data - mean_broad, data - mean_broad)
        axes = list(range(var.ndim))
        del axes[1]
        var = mx.np.mean(var, axis=axes)
        return mean, var

    def check_quantized_bn(data_shape, qdtype):
        if is_test_for_native_cpu():
            print('skipped testing quantize_bn for native cpu since it is not supported yet')
            return
        elif is_test_for_gpu():
            print('skipped testing quantize_bn for gpu since it is not supported yet')
            return

        data_low, data_high = get_low_high(qdtype)

        # run fp32 bn
        bn_fp32 = mx.gluon.nn.BatchNorm(use_global_stats=True, scale=True)
        data = mx.np.random.uniform(low=data_low, high=data_high, size=data_shape)
        bn_fp32.initialize()
        bn_fp32.hybridize()
        bn_fp32(data)
        fp32_params = bn_fp32.collect_params()

        data = mx.np.random.uniform(low=data_low, high=data_high, size=data_shape)
        gamma = mx.np.random.uniform(low=data_low, high=data_high, size=fp32_params['gamma'].shape)
        beta = mx.np.random.uniform(low=data_low, high=data_high, size=fp32_params['beta'].shape)
        running_mean, running_var = get_mean_var(data)
        new_params = {
            'gamma':gamma,
            'beta':beta,
            'running_mean': running_mean,
            'running_var': running_var
        }

        bn_fp32.load_dict(new_params)
        output = bn_fp32(data)

        # generate int8 bn from fp32 bn
        calib_data = mx.gluon.data.DataLoader(data, batch_size=data_shape[0])
        quant_bn = mx.contrib.quant.quantize_net(bn_fp32,
                                                 quantized_dtype=qdtype,
                                                 quantize_mode='full',
                                                 calib_data=calib_data,
                                                 calib_mode='naive',
                                                 num_calib_batches=1,
                                                 device=mx.current_device())

        output_int8_to_fp32 = quant_bn(data)

        assert_almost_equal(output.asnumpy(), output_int8_to_fp32.asnumpy(), rtol=1e-1, atol=8)

    for qdtype in ['int8', 'uint8']:
      check_quantized_bn((32, 512, 4, 4), qdtype)
      check_quantized_bn((32, 1024, 8, 8), qdtype)
      check_quantized_bn((32, 3, 224, 224), qdtype)


def test_quantized_reshape():
    test_cases = [((2, 3, 5, 5),  (-2, -1),         False, (2, 75)),
                  ((2, 3, 5, 5),  (-2, -2, -1),     False, (2, 3, 25)),
                  ((5, 3, 4, 5),  (-2, -1, -2),     False, (5, 15, 4)),
                  ((2, 3, 5, 4),  (-1, -2, -2),     False, (8, 3, 5)),
                  ((2, 3, 5, 5),  (-2, -2, -2, -2), False, (2, 3, 5, 5)),
                  ((2, 1, 4, 5),  (-2, -3, -2, -2), False, (2, 4, 5)),
                  ((1, 1, 4, 1),  (-3, -3, -2, -2), False, (4, 1)),
                  ((1, 1, 1, 1),  (-3, -3, -3, -3), False, ()),
                  ((2, 4, 5, 3),  (-1, 2, 2, 1),    False, (30, 2, 2, 1)),
                  ((2, 3, 5, 6),  (-4,),            False, (2, 3, 5, 6)),
                  ((2, 3, 5, 6),  (6, 1, -4),       False, (6, 1, 5, 6)),
                  ((2, 3, 5, 6),  (-5, -5),         False, (6, 30)),
                  ((2, 3, 5, 6),  (-5, -1),         False, (6, 30)),
                  ((64,),         (-6, 16, 4),      False, (16, 4)),
                  ((64,),         (-6, 16, -1),     False, (16, 4)),
                  ((64, 1, 2, 3), (-6, 16, -1, -4), False, (16, 4, 1, 2, 3)),
                  ((8, 5, 4, 6),  (-4, -1, 3, -6),  True,  (8, 5, 4, 2, 3))]

    def check_quantized_reshape(shape, qdtype, newshape, reverse, expected_ret_shape):
        data_low, data_high = get_low_high(qdtype)
        qdata = mx.np.random.uniform(low=data_low, high=data_high, size=shape).astype(qdtype)
        min_data = mx.np.array([-1023.343], dtype='float32')
        max_data = mx.np.array([2343.324275], dtype='float32')
        qoutput, min_output, max_output = npx.quantized_reshape(qdata, min_data, max_data, newshape=newshape, reverse=reverse)
        assert qoutput.shape == expected_ret_shape
        assert same(qdata.asnumpy().flatten(), qoutput.asnumpy().flatten())
        assert same(min_data.asnumpy(), min_output.asnumpy())
        assert same(max_data.asnumpy(), max_output.asnumpy())

    for qdtype in ['int8', 'uint8']:
        for shape, newshape, reverse, expected_ret_shape in test_cases:
            check_quantized_reshape(shape, qdtype, newshape, reverse, expected_ret_shape)


def test_quantize_params():
    if is_test_for_native_cpu():
        print('skipped testing quantized_params for native cpu since it is not supported yet')
        return

    data = mx.sym.Variable('data')
    conv = mx.sym.Convolution(data, kernel=(1, 1), num_filter=2048, name='conv')
    sym = mx.sym.BatchNorm(data=conv, eps=2e-05, fix_gamma=False, momentum=0.9, use_global_stats=False, name='bn')
    offline_params = [name for name in sym.list_arguments()
                      if not name.startswith('data') and not name.endswith('label')]
    params = {}
    for name in offline_params:
        params[name] = mx.nd.uniform(shape=(2, 2))
    qsym, _ = mx.contrib.quant._quantize_symbol(sym, device=mx.current_device(),
                                                offline_params=offline_params, quantize_mode='full')
    qparams = mx.contrib.quant._quantize_params(qsym, params, min_max_dict = {})
    param_names = params.keys()
    qparam_names = qparams.keys()
    for name in qparam_names:
        if name.startswith('bn'):
            assert name in param_names
        elif name.startswith('conv'):
            assert name not in param_names
            assert name.find('quantize') != -1


class FP32Net(mx.gluon.nn.HybridBlock):
    def __init__(self, **kwargs):
        super(FP32Net, self).__init__(**kwargs)
        self.conv = mx.gluon.nn.Conv2D(channels=16, kernel_size=(1,1))
        self.bn = mx.gluon.nn.BatchNorm(epsilon=2e-05, scale=True, momentum=0.9, use_global_stats=False)
        self.act = mx.gluon.nn.Activation(activation='relu')
        self.pool = mx.gluon.nn.AvgPool2D(pool_size=(4,4))
        self.fc = mx.gluon.nn.Dense(units=10, flatten=True)

    def forward(self, x):
        out = self.conv(x)
        out = self.bn(out)
        out = self.act(out)
        out = self.pool(out)
        out = self.fc(out)
        return npx.softmax(out)


class FP32MultipleOutputs(mx.gluon.nn.HybridBlock):
    def __init__(self, length, **kwargs):
        super(FP32MultipleOutputs, self).__init__(**kwargs)
        self.length = length
        self.convs = mx.gluon.nn.Conv2D(channels=16, kernel_size=(1,1))
        self.fc = mx.gluon.nn.Dense(units=10, flatten=True)

    def forward(self, x):
        res = npx.slice_channel(x, num_outputs=self.length,
                                axis=1, squeeze_axis=1)
        out = []
        for i in range(self.length):
            out.append(self.convs(res[i]))
            out[i] = mx.np.expand_dims(out[i], axis=0)
        out = mx.np.concatenate(out)
        out = mx.np.reshape(out, ((self.length, -1)))
        out = self.fc(out)
        return npx.softmax(out)

class FP32MultipleInputs(mx.gluon.nn.HybridBlock):
    def __init__(self, **kwargs):
        super(FP32MultipleInputs, self).__init__(**kwargs)
        self.conv1 = mx.gluon.nn.Conv2D(channels=64, kernel_size=(1,1), use_bias=False)
        self.bn1 = mx.gluon.nn.BatchNorm()
        self.conv2 = mx.gluon.nn.Conv2D(channels=64, kernel_size=(1,1), use_bias=False)
        self.bn2 = mx.gluon.nn.BatchNorm()

    def forward(self, data0, data1):
        out0 = self.conv1(data0)
        out0 = self.bn1(out0)
        out1 = self.conv2(data1)
        out1 = self.bn2(out1)
        return out1 + out0

@use_np
@xfail_when_nonstandard_decimal_separator
def test_quantize_model():
    def check_params(params, qparams, qsym=None):
        if qsym is None:
            assert len(params) == len(qparams)
            for k, v in params.items():
                assert k in qparams
                assert same(v.asnumpy(), qparams[k].asnumpy())
        else:
            qparams_ground_truth = mx.contrib.quant._quantize_params(qsym, params, min_max_dict = {})
            assert len(qparams) == len(qparams_ground_truth)
            for k, v in qparams_ground_truth.items():
                assert k in qparams
                assert same(v.asnumpy(), qparams[k].asnumpy())

    def check_qsym_calibrated(qsym):
        attrs = qsym.attr_dict()
        for k, v in attrs.items():
            if k.find('requantize_') != -1:
                assert 'min_calib_range' in v
                assert 'max_calib_range' in v

    def check_qsym_qdtype(qsym, qdtype):
        attrs = qsym.attr_dict()
        for k, v in attrs.items():
            if k.find('_quantize') != -1:
                assert 'out_type' in v
                assert v['out_type'] == qdtype

    def skip_not_supported():
        if is_test_for_native_cpu():
            print('skipped testing quantize_model for native cpu since it is not supported yet')
            return True
        elif qdtype == 'int8' and is_test_for_dnnl():
            print('skipped testing quantize_model for oneDNN cpu int8 since it is not supported yet')
            return True
        elif qdtype == 'uint8' and is_test_for_gpu():
            print('skipped testing quantize_model for gpu uint8 since it is not supported yet')
            return True
        return False

    def check_quantize_model(qdtype):
        if is_test_for_native_cpu():
            print('skipped testing quantize_model for native cpu since it is not supported yet')
            return
        elif qdtype == 'int8' and is_test_for_dnnl():
            print('skipped testing quantize_model for oneDNN cpu int8 since it is not supported yet')
            return
        elif qdtype == 'uint8' and is_test_for_gpu():
            print('skipped testing quantize_model for gpu uint8 since it is not supported yet')
            return

        standard_net = FP32Net()
        standard_net.initialize()
        batch_size = 4
        data_shape = (batch_size, 4, 10, 10)

        length = batch_size  # specify num of outputs from split op
        multi_out_net = FP32MultipleOutputs(length)
        multi_out_net.initialize()
        multi_out_data_shape = (length, 4, 4, 10, 10)

        for net, dshape in zip((standard_net, multi_out_net), (data_shape, multi_out_data_shape)):
            data = mx.np.random.uniform(low=0, high=1, size=dshape)
            net.hybridize()
            net(data)
            sym, _ = net.export(None)
            arg_params, aux_params = collect_block_args_aux(net, sym)

            qsym, qarg_params, qaux_params = mx.contrib.quant.quantize_model(sym=sym,
                                                                             arg_params=arg_params,
                                                                             aux_params=aux_params,
                                                                             device=mx.current_device(),
                                                                             quantized_dtype=qdtype,
                                                                             calib_mode='none',
                                                                             quantize_mode='full')
            check_params(arg_params, qarg_params, qsym)
            check_params(aux_params, qaux_params)

            calib_data = mx.np.random.uniform(size=dshape)
            calib_data = mx.gluon.data.DataLoader(calib_data, batch_size=batch_size)
            qsym, qarg_params, qaux_params = mx.contrib.quant.quantize_model(sym=sym,
                                                                             arg_params=arg_params,
                                                                             aux_params=aux_params,
                                                                             device=mx.current_device(),
                                                                             quantized_dtype=qdtype,
                                                                             calib_mode='naive',
                                                                             calib_data=calib_data,
                                                                             num_calib_batches=1,
                                                                             quantize_mode='full')
            check_params(arg_params, qarg_params, qsym)
            check_params(aux_params, qaux_params)
            check_qsym_calibrated(qsym)
            check_qsym_qdtype(qsym, qdtype)

    def check_quantize_model_multiple_inputs(qdtype):
        if skip_not_supported():
            return

        net = FP32MultipleInputs()
        net.initialize()
        net.hybridize()
        dshape = (64, 4, 10, 10)
        data = [mx.np.random.uniform(low=0, high=1, size=dshape),
                mx.np.random.uniform(low=0, high=1, size=dshape)]
        net(*data)
        sym, _ = net.export(None)
        arg_params, aux_params = collect_block_args_aux(net, sym)

        qsym, qarg_params, qaux_params = mx.contrib.quant.quantize_model(sym=sym,
                                                                         arg_params=arg_params,
                                                                         aux_params=aux_params,
                                                                         device=mx.current_device(),
                                                                         quantized_dtype=qdtype,
                                                                         calib_mode='none',
                                                                         quantize_mode='full')
        check_params(arg_params, qarg_params, qsym)
        check_params(aux_params, qaux_params)

        calib_data = [mx.np.random.uniform(size=dshape),
                      mx.np.random.uniform(size=dshape)]
        calib_data = mx.gluon.data.DataLoader(mx.gluon.data.ArrayDataset(*calib_data), batch_size=4)
        qsym, qarg_params, qaux_params = mx.contrib.quant.quantize_model(sym=sym,
                                                                         arg_params=arg_params,
                                                                         aux_params=aux_params,
                                                                         device=mx.current_device(),
                                                                         quantized_dtype=qdtype,
                                                                         calib_mode='naive',
                                                                         calib_data=calib_data,
                                                                         data_names=["data0","data1"],
                                                                         num_calib_batches=1,
                                                                         quantize_mode='full')
        check_params(arg_params, qarg_params, qsym)
        check_params(aux_params, qaux_params)
        check_qsym_calibrated(qsym)
        check_qsym_qdtype(qsym, qdtype)


    for qdtype in ['int8', 'uint8']:
        check_quantize_model(qdtype)
        check_quantize_model_multiple_inputs(qdtype)


@mx.util.use_np
def test_quantize_gluon_with_forward():
    def check_quantize_net(qdtype):
        if is_test_for_native_cpu():
            print('skipped testing test_quantize_model_with_forward for native cpu since it is not supported yet')
            return
        elif is_test_for_gpu():
            print('skipped testing test_quantize_model_with_forward for gpu uint8 since it is not supported yet')
            return

        data_shape = (32, 3, 224, 224)
        batch_size = 1
        resnet18_v1 = vision.resnet18_v1(pretrained=True)
        resnet18_v1.reset_device(mx.current_device())
        excluded_names_match = []
        if mx.current_device() == mx.gpu():
            excluded_names_match += ['activation', 'relu', 'conv0']
        num_calib_batches = 1

        random_data = mx.np.random.uniform(size=data_shape)
        calib_data = mx.gluon.data.DataLoader(random_data, batch_size=batch_size)

        quantized_resnet18_v1 = mx.contrib.quant.quantize_net(resnet18_v1, quantized_dtype=qdtype,
                                                              exclude_layers=None,
                                                              exclude_layers_match=excluded_names_match,
                                                              calib_mode='none',
                                                              data_shapes=[data_shape],
                                                              device=mx.current_device())
        quantized_resnet18_v1.hybridize(static_alloc=True, static_shape=True)
        quantized_resnet18_v1(random_data)

        for mode in ['naive', 'entropy']:
            for quantize_granularity in ['tensor-wise', 'channel-wise']:
                qdtype = qdtype if mode == 'naive' else 'auto'
                quantized_resnet18_v1 = mx.contrib.quant.quantize_net(resnet18_v1, quantized_dtype=qdtype,
                                                                    exclude_layers=None,
                                                                    exclude_layers_match=excluded_names_match,
                                                                    calib_data=calib_data,
                                                                    calib_mode=mode,
                                                                    quantize_granularity=quantize_granularity,
                                                                    num_calib_batches=num_calib_batches,
                                                                    device=mx.current_device())

                quantized_resnet18_v1.hybridize(static_alloc=True, static_shape=True)
                quantized_resnet18_v1(random_data)

    for qdtype in ['int8', 'uint8']:
        check_quantize_net(qdtype)


@xfail_when_nonstandard_decimal_separator
def test_quantize_sym_with_calib():
    if is_test_for_native_cpu():
        print('skipped testing quantized_pooling for native cpu since it is not supported yet')
        return

    def get_fp32_sym():
        data = mx.sym.Variable('data')
        conv = mx.sym.Convolution(data, kernel=(1, 1), num_filter=16, name='conv')
        bn = mx.sym.BatchNorm(data=conv, eps=2e-05, fix_gamma=False, momentum=0.9, use_global_stats=False, name='bn')
        act = mx.sym.Activation(data=bn, act_type='relu', name='relu')
        pool = mx.sym.Pooling(act, kernel=(4, 4), pool_type='avg', name='pool')
        fc = mx.sym.FullyConnected(pool, num_hidden=10, flatten=True, name='fc')
        sym = mx.sym.softmax(fc, name='softmax')
        return sym

    sym = get_fp32_sym()
    offline_params = [name for name in sym.list_arguments()
                      if not name.startswith('data') and not name.endswith('label')]
    qsym, _ = mx.contrib.quant._quantize_symbol(sym, device=mx.current_device(),
                                             offline_params=offline_params, quantize_mode='full')
    requantize_op_names = ['requantize_conv', 'requantize_fc']
    min_max_dict = {'conv_output': (onp.random.uniform(low=100.0, high=200.0), onp.random.uniform(low=100.0, high=200.0)),
                    'fc_output': (onp.random.uniform(low=100.0, high=200.0), onp.random.uniform(low=100.0, high=200.0))}
    op_name_to_th_name = {'requantize_conv': 'conv_output', 'requantize_fc': 'fc_output'}
    cqsym = mx.contrib.quant._calibrate_quantized_sym(qsym, min_max_dict)
    attr_dict = cqsym.attr_dict()
    for name in requantize_op_names:
        assert name in attr_dict
        lhs = float(attr_dict[name]['min_calib_range'])
        rhs = min_max_dict[op_name_to_th_name[name]][0]
        assert_almost_equal(onp.array([lhs]), onp.array([rhs]))
        lhs = float(attr_dict[name]['max_calib_range'])
        rhs = min_max_dict[op_name_to_th_name[name]][1]
        assert_almost_equal(onp.array([lhs]), onp.array([rhs]), rtol=1e-3, atol=1e-4)


@use_np
def test_quantization_net_with_different_data_inputs_options():
    if is_test_for_native_cpu():
        print('skipped testing test_quantization_net_with_different_data_inputs_options for native cpu since it is not supported yet')
        return
    elif is_test_for_gpu():
        print('skipped testing test_quantization_net_with_different_data_inputs_options for gpu since it is not supported yet')
        return

    net = FP32Net()
    net.initialize()

    batch_size = 32
    data_shape = (batch_size, 3, 224, 224)
    random_data = mx.np.random.uniform(size=data_shape)

    # pass data_shapes as list of tuples
    quantized_net = mx.contrib.quant.quantize_net(net,
                                                  quantized_dtype='auto',
                                                  data_shapes=[data_shape],
                                                  device=mx.current_device())
    out = quantized_net(random_data)
    out.wait_to_read()


    # pass data_shapes as list of DataDescs
    net2 = FP32Net()
    net2.initialize()
    data_desc = mx.io.DataDesc('data', data_shape)
    quantized_net2 = mx.contrib.quant.quantize_net(net2,
                                                   quantized_dtype='auto',
                                                   data_shapes=[data_desc],
                                                   device=mx.current_device())
    out2 = quantized_net2(random_data)
    out2.wait_to_read()


    # pass data as DataLoader
    net3 = FP32Net()
    net3.initialize()
    data_loader = mx.gluon.data.DataLoader(random_data, batch_size=batch_size)
    quantized_net3 = mx.contrib.quant.quantize_net(net3,
                                                   quantized_dtype='auto',
                                                   calib_data=data_loader,
                                                   device=mx.current_device())
    out3 = quantized_net3(random_data)
    out3.wait_to_read()


def test_optimal_threshold_adversarial_case():
    # The worst case for the optimal_threshold function is when the values are concentrated
    # at one edge: [0, 0, ..., 1000]. (histogram)
    # We want to make sure that the optimal threshold in this case is the max.
    hist = []
    hist_edges = []
    min_val = -2
    max_val = 2
    for _ in range(0, 998):
        hist.append(0)
    for i in range(0, 999):
        hist_edges.append((max_val - min_val) / 999 * i + min_val)
    hist.append(1000)
    hist_edges.append(max_val)
    hist_data = (hist, hist_edges, min_val, max_val, max_val)
    for dtype in ['uint8', 'int8', 'auto']:
        res = mx.contrib.quant._LayerHistogramCollector.get_optimal_threshold(hist_data, dtype, num_quantized_bins=5)
        # The threshold should be 2.
        print (res)
        assert abs(res[2] - 2) < 1e-5


def test_get_optimal_thresholds():
    # Given an ndarray with elements following a uniform distribution, the optimal threshold
    # for quantizing the ndarray should be either abs(min(nd)) or abs(max(nd)).
    def get_threshold(nd):
        min_nd = mx.nd.min(nd)
        max_nd = mx.nd.max(nd)
        return mx.nd.maximum(mx.nd.abs(min_nd), mx.nd.abs(max_nd)).asnumpy()

    for dtype in ['uint8', 'int8', 'auto']:
        nd = mx.nd.uniform(low=-10.532, high=11.3432, shape=(8, 3, 23, 23), dtype=onp.float64)
        expected_threshold = get_threshold(nd)
        arr = nd.asnumpy()
        min_range = onp.min(arr)
        max_range = onp.max(arr)
        th = max(abs(min_range), abs(max_range))
        hist, hist_edges = onp.histogram(arr, bins=8001, range=(-th, th))
        hist_dict = {'layer1' : (hist, hist_edges, min_range, max_range, th)}
        min_max_dict = mx.contrib.quant._LayerHistogramCollector.get_optimal_thresholds(hist_dict, dtype)
        assert 'layer1' in min_max_dict
        assert_almost_equal(onp.array([min_max_dict['layer1'][1]]), expected_threshold, rtol=1e-2, atol=1e-4)


@use_np
def test_rnn_quantization():
    data_low = -1
    data_high = 1
    def check_rnn_quantization(num_layers, bidirectional, seq_len, batch_size, input_dim, state_size):
        data_shape = (seq_len, batch_size, input_dim)

        rnn_fp32 = mx.gluon.rnn.LSTM(hidden_size=state_size,
                                     num_layers = num_layers,
                                     bidirectional=bidirectional)

        data = mx.np.random.uniform(low=data_low, high=data_high, size=data_shape)
        states_shape = (num_layers * 2 if bidirectional else num_layers, batch_size, state_size)
        states = [mx.np.zeros((states_shape)) for _ in range(batch_size)]

        rnn_fp32.initialize()
        rnn_fp32.hybridize()
        ref_out = rnn_fp32(data, states)

        class RNNDataLoader(mx.gluon.data.DataLoader):
            def __init__(self, data, states):
                super().__init__(mx.gluon.data.SimpleDataset([]), 1)
                self.data = data
                self.states = states

            def __iter__(self):
                return self

            def __next__(self):
                return [self.data, self.states]

            def __bool__(self):
                return bool(self.dataiter.iter_next())

        calib_data = RNNDataLoader(data, states)
        quant_rnn = mx.contrib.quant.quantize_net(rnn_fp32,
                                                  quantized_dtype='auto',
                                                  quantize_mode='full',
                                                  calib_data=calib_data,
                                                  calib_mode='naive',
                                                  num_calib_batches=1,
                                                  device=mx.current_device())
        qout = quant_rnn(data, states)

        qsym, _ = quant_rnn.export(None)
        assert qsym.tojson().find("quantized_rnn") != -1

        ref_out = [ref_out[0], ref_out[1][0], ref_out[1][1]]
        for i in range(len(qout)):
            mse = onp.mean((ref_out[i].asnumpy() - qout[i].asnumpy())**2)
            assert mse < 0.001

    check_rnn_quantization(1, False, 5, 2, 16, 16)
    check_rnn_quantization(1, True, 5, 2, 16, 16)


@use_np
def test_quantized_rnn():
    def check_quantized_rnn(num_layers, bidirectional, seq_len, batch_size, input_dim, state_size):
        ndir = 2 if bidirectional else 1
        size = ndir*state_size*4
        first_lyr_param_size = (input_dim + state_size + 2) * size
        other_lyr_param_size = (state_size * ndir + state_size + 2) * size
        full_param_size = first_lyr_param_size + (num_layers - 1) * other_lyr_param_size

        data = mx.np.random.uniform(-1, 1, (seq_len, batch_size, input_dim))
        state = mx.np.random.uniform(-1, 1, (num_layers*ndir, batch_size, state_size))
        state_cell = mx.np.random.uniform(0, 1, (num_layers*ndir, batch_size, state_size))
        params = mx.np.random.normal(0, 1, (full_param_size,))

        out = npx.rnn(data=data,
                      parameters=params,
                      mode='lstm',
                      state=state,
                      state_size=state_size,
                      state_cell=state_cell,
                      num_layers=num_layers,
                      bidirectional=bidirectional)

        data_min = mx.np.min(data)
        data_max = mx.np.max(data)
        data_scale = mx.np.array(128.0 / (data_max - data_min)).reshape((1,))
        data_shift = mx.np.array(128.0 - data_max * data_scale).reshape((1,))

        qdata = (data * data_scale + data_shift + 0.5).astype('uint8')
        qout = npx.contrib_quantized_rnn(data=qdata,
                                         parameters=params,
                                         mode='lstm',
                                         state=state,
                                         state_size=state_size,
                                         state_cell=state_cell,
                                         num_layers=num_layers,
                                         bidirectional=bidirectional,
                                         data_scale=data_scale,
                                         data_shift=data_shift)

        mse = onp.mean((out.asnumpy() - qout.asnumpy())**2)
        assert mse < 0.001

    check_quantized_rnn(1, False, 5, 2, 16, 16)
    check_quantized_rnn(1, True, 5, 2, 16, 16)