#
# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed 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.
#
"""tests of tensor quantization function and module"""
import contextlib

import pytest
import numpy as np

import torch
from torch.nn.parameter import Parameter

from pytorch_quantization import cuda_ext
from pytorch_quantization import tensor_quant
import pytorch_quantization.utils as quant_utils

import tests.utils as test_utils
from tests.fixtures import verbose

np.random.seed(123456)  # seed 1234 causes 1 number mismatch at 6th decimal in one of the tests

# pylint:disable=missing-docstring, no-self-use


class TestTensorQuant():

    def test_simple_run(self):
        """ quantizer passes gradcheck
        """
        x = Parameter(torch.randn(2, 3, dtype=torch.float64).cuda()) * 100
        tensor_quant.tensor_quant(x, torch.max(torch.abs(x)), 7)

    def test_per_tensor_scale(self):
        """ tensor_quant matches numpy quantization
        """
        torch.set_default_tensor_type('torch.cuda.FloatTensor')  # Test on GPU
        x_np = np.random.rand(1023)
        x_torch = torch.Tensor(x_np)
        quant_x_np = test_utils.quant_np(x_np, np.max(np.abs(x_np)))
        quant_x_torch, _ = tensor_quant.tensor_quant(x_torch, torch.max(torch.abs(x_torch)))
        np.testing.assert_array_equal(quant_x_torch.cpu().numpy(), quant_x_np)
        torch.set_default_tensor_type('torch.FloatTensor')

    def test_per_channel_scale(self):
        """ fake_tensor_quant performs per channel quantization
        """
        x_np = np.random.rand(15, 15, 64, 128).astype('float32')
        x_torch = torch.Tensor(x_np).cuda()

        # Pytorch filter layout seems to be KCRS, reduce max to shape [K, 1, 1, 1] to test per channel scale
        # Shrink max a little, so that clip behavior is tested
        amax_x_np = 0.7 * np.max(np.abs(x_np), axis=(1, 2, 3), keepdims=True)
        # Pytorch's max function doesn't support reduces multiple axis, and returns (max, argmax) tuple,
        # so it has to be reduced by multiple torch.max
        amax_x_torch = 0.7 * torch.max(
            torch.max(torch.max(x_torch, dim=1, keepdim=True)[0], dim=2, keepdim=True)[0], dim=3, keepdim=True)[0]

        quant_x_np = test_utils.quant_np(x_np, amax_x_np)
        quant_x_torch, _ = tensor_quant.tensor_quant(x_torch, amax_x_torch)

        # np.testing.assert_array_equal(quant_x_torch.cpu().numpy(), quant_x_np)
        # Pytorch numerics is not the same as numpy, it will be off by 1
        np.testing.assert_array_less(np.abs(quant_x_torch.cpu().numpy() - quant_x_np), 2)
        if verbose:
            mismatches = np.where(np.abs(quant_x_torch.cpu().numpy() - quant_x_np) >= 1)
            print("Mismatches:")
            print(" Original: ", x_np[mismatches])
            print(" numpy: ", quant_x_np[mismatches])
            print(" Pytorch: ", quant_x_torch.cpu().numpy()[mismatches])

    def test_backward(self):
        """ tensor_quant implements straight through estimator on the backward pass
            Note: this does not work for integer output_dtype
        """
        x = torch.randn(3, 7, requires_grad=True).cuda()
        labels = torch.randint(6, (3,)).type(torch.LongTensor).cuda()
        quant_x, _ = tensor_quant.tensor_quant(x, x.abs().max(), 7)
        float_quant_x = quant_x.type(torch.FloatTensor).cuda()
        x.retain_grad()
        float_quant_x.retain_grad()
        criterion = torch.nn.CrossEntropyLoss().cuda()
        loss = criterion(float_quant_x, labels)
        loss.backward()
        np.testing.assert_array_equal(float_quant_x.grad.cpu().numpy(), x.grad.cpu().numpy())

    def test_unsigned(self):
        x_np = np.random.rand(1023).astype('float32')
        x_torch = torch.Tensor(x_np)
        quant_x_np = test_utils.quant_np(x_np, np.max(np.abs(x_np)), num_bits=9, fake=False)
        quant_x_torch, _ = tensor_quant.tensor_quant(x_torch, torch.max(torch.abs(x_torch)), 8, True)
        np.testing.assert_array_almost_equal(quant_x_torch.cpu().numpy(), quant_x_np)

        x_torch = torch.randn(3, 7)
        with pytest.raises(TypeError, match="Negative values encountered"):
            tensor_quant.tensor_quant(x_torch, torch.max(torch.abs(x_torch)), 8, True)

    def test_overflow_fp16(self):
        x_torch = torch.randn(1023).cuda().half()
        with pytest.raises(ValueError, match="scale is too large for FP16"):
            quant_x_torch, scale = tensor_quant.tensor_quant(x_torch, torch.tensor(1e-4).cuda().half(), 8, False)

    def test_clip_gradient(self):
        x = torch.randn(3, 7, requires_grad=True).cuda()
        x.retain_grad()
        amax = x.abs().max() / 2
        x_in_range = (-amax <= x) * (x <= amax)
        quant_x, _ = tensor_quant.tensor_quant(x, amax, 8)
        loss = torch.sum((quant_x - 0.5)**2)
        loss.backward()
        np.testing.assert_array_equal(x.grad.cpu().numpy() != 0, x_in_range.cpu().numpy())

    def test_full_range(self):
        """ fake_tensor_quant uses the full integer range when narrow=False
        """
        x_np = np.random.rand(1023).astype('float32')
        x_torch = torch.Tensor(x_np).cuda()
        amax = np.max(np.abs(x_np))
        quant_x_np = test_utils.quant_np(x_np, amax, num_bits=9, fake=False, narrow_range=False)
        quant_x_torch, _ = tensor_quant.tensor_quant(x_torch, torch.max(torch.abs(x_torch)), 8, True, False)
        np.testing.assert_array_almost_equal(quant_x_torch.cpu().numpy(), quant_x_np)


class TestFakeTensorQuant():

    def test_simple_run(self):
        x = Parameter(torch.randn(3, 7).cuda())
        tensor_quant.fake_tensor_quant(x, torch.max(torch.abs(x)))

    def test_per_tensor_scale(self):
        """ fake_tensor_quant matches numpy quantization
        """
        x_np = np.random.rand(13).astype('float32')
        print(x_np)
        x_torch = torch.Tensor(x_np).cuda()
        quant_x_np = test_utils.quant_np(x_np, np.max(np.abs(x_np)), fake=True)
        quant_x_torch = tensor_quant.fake_tensor_quant(x_torch, torch.max(torch.abs(x_torch)))
        np.testing.assert_array_almost_equal(quant_x_torch.cpu().numpy(), quant_x_np)

    def test_per_channel_scale(self):
        """ fake_tensor_quant performs per channel quantization
        """
        x_np = np.random.rand(15, 15, 64, 128).astype('float32')
        x_torch = torch.Tensor(x_np).cuda()

        # Pytorch filter layout seems to be KCRS, reduce max to shape [K, 1, 1, 1] to test per channel scale
        # Shrink max a little, so that clip behavior is tested
        amax_x_np = 0.9 * np.max(np.abs(x_np), axis=(1, 2, 3), keepdims=True)
        # Pytorch's max function doesn't support reduces multiple axis, and returns (max, argmax) tuple,
        # so it has to be reduced by multiple torch.max
        amax_x_torch = 0.9 * torch.max(
            torch.max(torch.max(x_torch, dim=1, keepdim=True)[0], dim=2, keepdim=True)[0], dim=3, keepdim=True)[0]

        quant_x_np = test_utils.quant_np(x_np, amax_x_np, fake=True)
        quant_x_torch = tensor_quant.fake_tensor_quant(x_torch, amax_x_torch)

        # Pytorch numerics is not the same as numpy, results will be off a little
        # np.testing.assert_array_equal(quant_x_torch.cpu().numpy(), quant_x_np)
        np.testing.assert_array_almost_equal(quant_x_torch.cpu().numpy(), quant_x_np, decimal=2)
        if verbose:
            mismatches = np.where(np.abs(quant_x_torch.cpu().numpy() - quant_x_np) >= 1e-5)
            print("Mismatches:")
            print(" Original: ", x_np[mismatches])
            print(" numpy: ", quant_x_np[mismatches])
            print(" Pytorch: ", quant_x_torch.cpu().numpy()[mismatches])

    def test_backward(self):
        """ fake_tensor_quant implements straight through estimator on the backward pass
        """
        x = torch.randn(3, 7, requires_grad=True).cuda()
        labels = torch.randint(6, (3,)).type(torch.LongTensor).cuda()
        quant_x = tensor_quant.fake_tensor_quant(x, torch.max(torch.abs(x)), 7)
        x.retain_grad()
        quant_x.retain_grad()
        criterion = torch.nn.CrossEntropyLoss().cuda()
        loss = criterion(quant_x, labels)
        loss.backward()
        np.testing.assert_array_equal(quant_x.grad.cpu().numpy(), x.grad.cpu().numpy())

    def test_unsigned(self):
        x_np = np.random.rand(1023).astype('float32')
        x_torch = torch.Tensor(x_np).cuda()
        quant_x_np = test_utils.quant_np(x_np, np.max(np.abs(x_np)), num_bits=9, fake=True)
        quant_x_torch = tensor_quant.fake_tensor_quant(x_torch, torch.max(torch.abs(x_torch)), 8, True)
        np.testing.assert_array_almost_equal(quant_x_torch.cpu().numpy(), quant_x_np)

    def test_cuda_ext(self):
        x_np = np.random.rand(1023).astype('float32')
        x_torch = torch.Tensor(x_np).cuda()

        for num_bits in [3, 4, 5, 7, 8, 11]:
            for unsigned in [True, False]:
                test_utils.compare(cuda_ext.fake_tensor_quant(x_torch, torch.max(torch.abs(x_torch)), num_bits,
                                                              unsigned),
                                   tensor_quant.fake_tensor_quant(x_torch, torch.max(torch.abs(x_torch)), num_bits,
                                                                  unsigned),
                                   rtol=0,
                                   atol=0)

        # Test fp16 and bf16
        for dtype in [torch.float16, torch.bfloat16]:
            x_np = np.random.rand(1023)
            x_torch = torch.Tensor(x_np).cuda().to(dtype)
            cuda_ext_out = cuda_ext.fake_tensor_quant(x_torch, torch.max(torch.abs(x_torch))).to(torch.float32)
            pytorch_out = tensor_quant.fake_tensor_quant(x_torch, torch.max(torch.abs(x_torch))).to(torch.float32)
            test_utils.compare(cuda_ext_out, pytorch_out, rtol=0, atol=0)

    @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
    def test_cuda_ext_with_axis(self, dtype):
        x_np = np.random.rand(3, 4, 5, 6)
        x_torch = torch.Tensor(x_np).cuda().to(dtype)

        # amax along axis 1
        amax_torch = torch.tensor([0.8, 0.9, 0.7, 0.6], device="cuda")

        for num_bits in [3, 4, 5, 7, 8, 11]:
            for unsigned in [True, False]:
                cuda_ext_out = cuda_ext.fake_tensor_quant_with_axis(x_torch, amax_torch, 1, num_bits, unsigned).to(torch.float32)
                pytorch_out = tensor_quant.fake_tensor_quant(x_torch, amax_torch.view(1, -1, 1, 1), num_bits, unsigned).to(torch.float32)
                test_utils.compare(cuda_ext_out, pytorch_out, rtol=0, atol=0)

    def test_cuda_ext_inplace(self):
        x_np = np.random.rand(1023).astype('float32')
        x_torch = torch.Tensor(x_np).cuda()
        quant_x_np = test_utils.quant_np(x_np, np.max(np.abs(x_np)), fake=True)
        cuda_ext.fake_tensor_quant_(x_torch, torch.max(torch.abs(x_torch)))
        np.testing.assert_array_equal(x_torch.cpu().numpy(), quant_x_np)

        # Test fp16 and bf16
        for dtype in [torch.float16, torch.bfloat16]:
            x_np = np.random.rand(1023)
            x_torch = torch.Tensor(x_np).cuda().to(dtype)
            quant_x_np = test_utils.quant_np(x_np, np.max(np.abs(x_np)), fake=True)
            cuda_ext.fake_tensor_quant_(x_torch, torch.max(torch.abs(x_torch)))
            x_torch = x_torch.to(torch.float32)
            np.testing.assert_array_almost_equal(x_torch.cpu().numpy(), quant_x_np, decimal=2)

    def test_cuda_ext_tiny_amax(self):
        x_torch = torch.rand(2, 3, 4, device="cuda")
        amax = torch.tensor([1., 1.e-26, 1.], device="cuda").unsqueeze(-1).unsqueeze(1)
        quant_x = cuda_ext.fake_tensor_quant_with_axis(x_torch, amax, axis=1)
        assert quant_x[:, 1, :].sum() == 0

    def test_overflow_fp16(self):
        x_torch = torch.randn(1023).cuda().half()
        quant_x_torch = tensor_quant.fake_tensor_quant(x_torch, torch.tensor(1e-4).cuda().half(), 8, False)
        assert not (torch.isinf(quant_x_torch).any() or torch.isnan(quant_x_torch).any())

    def test_clip_gradient(self):
        x = torch.randn(3, 7, requires_grad=True).cuda()
        x.retain_grad()
        amax = x.abs().max() / 2
        x_in_range = (-amax <= x) * (x <= amax)
        quant_x = tensor_quant.fake_tensor_quant(x, amax, 8)
        loss = torch.sum((quant_x - 0.5)**2)
        loss.backward()
        np.testing.assert_array_equal(x.grad.cpu().numpy() != 0, x_in_range.cpu().numpy())

    def test_full_range(self):
        """ fake_tensor_quant uses the full integer range when narrow=False
        """
        x_np = np.random.rand(1023).astype('float32')
        x_torch = torch.Tensor(x_np).cuda()
        amax = np.max(np.abs(x_np))
        quant_x_np = test_utils.quant_np(x_np, amax, num_bits=9, fake=True, narrow_range=False)
        quant_x_torch = tensor_quant.fake_tensor_quant(x_torch, torch.max(torch.abs(x_torch)), 8, True, False)
        np.testing.assert_array_almost_equal(quant_x_torch.cpu().numpy(), quant_x_np)

    @pytest.mark.parametrize("dtype", ["float32", "float16"])
    def test_against_legacy(self, dtype):
        x_np = np.random.rand(3, 4, 5, 6).astype(dtype)
        x_torch = torch.Tensor(x_np).cuda()

        amax_torch = torch.tensor(0.7, device="cuda")

        for num_bits in [3, 4, 5, 7, 8, 11]:
            for unsigned in [True, False]:
                legacy_out = tensor_quant.legacy_fake_tensor_quant(x_torch, amax_torch, num_bits, unsigned)
                test_out = tensor_quant.fake_tensor_quant(x_torch, amax_torch, num_bits, unsigned)
                test_utils.compare(legacy_out, test_out, rtol=0, atol=0)

    def test_against_legacy_noncontiguous(self):
        x_np = np.random.rand(3, 4, 5, 6)
        x_torch = torch.Tensor(x_np).cuda()

        amax_torch = torch.tensor(0.7, device="cuda")

        x_torch_noncontiguous = x_torch[:, 2, :, 3]
        assert not x_torch_noncontiguous.is_contiguous()

        legacy_out = tensor_quant.legacy_fake_tensor_quant(x_torch_noncontiguous, amax_torch)
        test_out = tensor_quant.fake_tensor_quant(x_torch_noncontiguous, amax_torch)
        test_utils.compare(legacy_out, test_out, rtol=0, atol=0)

    @pytest.mark.parametrize("dtype", ["float32", "float16"])
    def test_against_legacy_with_axis(self, dtype):
        x_np = np.random.rand(3, 4, 5, 6).astype(dtype)
        x_torch = torch.Tensor(x_np).cuda()

        # amax along axis 1
        amax_torch = torch.tensor([0.8, 0.9, 0.7, 0.6], device="cuda").view(1, -1, 1, 1)

        for num_bits in [3, 4, 5, 7, 8, 11]:
            for unsigned in [True, False]:
                legacy_out = tensor_quant.legacy_fake_tensor_quant(x_torch, amax_torch, num_bits, unsigned)
                test_out = tensor_quant.fake_tensor_quant(x_torch, amax_torch, num_bits, unsigned)
                test_utils.compare(legacy_out, test_out, rtol=0, atol=0)


class TestQuantDescriptor():

    def test_scaled_mode(self):
        num_bits = np.random.randint(0, 16)

        test_quant_desc = tensor_quant.QuantDescriptor(num_bits=num_bits)
        assert test_quant_desc.num_bits == num_bits
        assert test_quant_desc.axis is None
        assert test_quant_desc.amax is None
        assert not test_quant_desc.learn_amax

        axis = (0, 1, 3)
        test_quant_desc = tensor_quant.QuantDescriptor(axis=axis)
        assert test_quant_desc.num_bits == 8  # default value
        assert test_quant_desc.axis == axis
        assert test_quant_desc.amax is None

        amax = 0.7
        test_quant_desc = tensor_quant.QuantDescriptor(amax=amax, unsigned=True)
        assert test_quant_desc.axis is None
        assert test_quant_desc.amax == np.float32(amax)
        assert test_quant_desc.unsigned

        amax = 0.7
        test_quant_desc = tensor_quant.QuantDescriptor(amax=amax, learn_amax=True)
        assert test_quant_desc.amax == np.float32(amax)
        assert test_quant_desc.learn_amax

        # Test the print string once if verbose is set.
        if verbose:
            print(test_quant_desc)

        with pytest.raises(TypeError, match="must be float, list or ndarray"):
            tensor_quant.QuantDescriptor(amax='oops')

        with pytest.raises(TypeError, match="amax must be float, list or ndarray"):
            tensor_quant.QuantDescriptor(amax='oops', learn_amax=True)

        with pytest.raises(TypeError, match="axis is ignored and must be None"):
            tensor_quant.QuantDescriptor(axis=(1, 2), amax=0.7, learn_amax=True)

    def test_amax(self):
        test_quant_desc = tensor_quant.QuantDescriptor()
        assert test_quant_desc.amax is None

        test_quant_desc = tensor_quant.QuantDescriptor(amax=1.2)
        assert isinstance(test_quant_desc.amax, np.ndarray)
        np.testing.assert_array_equal(test_quant_desc.amax, np.float32(1.2))

        test_quant_desc = tensor_quant.QuantDescriptor(amax=[1.3, 1.4])
        assert isinstance(test_quant_desc.amax, np.ndarray)
        np.testing.assert_array_equal(test_quant_desc.amax, np.float32([1.3, 1.4]))

        with pytest.raises(TypeError, match="must be float, list or ndarray"):
            tensor_quant.QuantDescriptor(amax='oops')

    def test_from_to_dict(self):
        quant_desc_1 = tensor_quant.QuantDescriptor(num_bits=2,
                                                    name='a',
                                                    fake_quant=True,
                                                    axis=(1, 2),
                                                    amax=3.1415926536)
        quant_desc_2 = tensor_quant.QuantDescriptor(**quant_desc_1.dict())
        if verbose:
            print(quant_desc_1.dict())
        assert quant_desc_1 == quant_desc_2

        quant_desc_1 = tensor_quant.QuantDescriptor(num_bits=2, amax=0.1, unsigned=True)
        quant_desc_2 = tensor_quant.QuantDescriptor(**quant_desc_1.dict())
        assert quant_desc_1 == quant_desc_2

    def test_from_to_yaml(self):
        quant_desc_1 = tensor_quant.QuantDescriptor(num_bits=2,
                                                    name='a',
                                                    fake_quant=True,
                                                    axis=(1, 2),
                                                    amax=3.1415926536)
        quant_desc_2 = tensor_quant.QuantDescriptor.from_yaml(quant_desc_1.to_yaml())
        if verbose:
            print(quant_desc_1.to_yaml())
        assert quant_desc_1 == quant_desc_2

        quant_desc_1 = tensor_quant.QuantDescriptor(num_bits=2, amax=0.1)
        quant_desc_2 = tensor_quant.QuantDescriptor.from_yaml(quant_desc_1.to_yaml())
        assert quant_desc_1 == quant_desc_2


class TestFakeAffineTensorQuant():

    def test_simple_run(self, verbose):
        x = np.array([-1., -13., -101., -128., 0., 2., 5., 13., 93., 111., 127.], dtype=np.float32)
        torch_x = torch.tensor(x).cuda()
        quant_x = tensor_quant.fake_affine_tensor_quant(torch_x, torch.min(torch_x), torch.max(torch_x))

        if verbose:
            print(quant_x)

        np.testing.assert_array_almost_equal(quant_x.cpu().numpy(), x)

    def test_clip_gradient(self):
        x = torch.randn(3, 7, requires_grad=True).cuda()
        x.retain_grad()
        xmin = x.min() / 2
        xmax = x.max() / 2
        x_in_range = (xmin <= x) * (x <= xmax)
        quant_x = tensor_quant.fake_affine_tensor_quant(x, xmin, xmax, 8)
        loss = torch.sum((quant_x - 0.5)**2)
        loss.backward()
        np.testing.assert_array_equal(x.grad.cpu().numpy() != 0, x_in_range.cpu().numpy())


class TestScaledE4M3():

    x = [[-2.0000, -1.8000, -1.6000, -1.4000, -1.2000], [-1.0000, -0.8000, -0.6000, -0.4000, -0.2000],
         [-0.0000, 0.2000, 0.4000, 0.6000, 0.8000], [1.0000, 1.2000, 1.4000, 1.6000, 1.8000]]

    xq_unscaled = [[-2.0000, -1.7500, -1.6250, -1.3750, -1.2500], [-1.0000, -0.8125, -0.6250, -0.4062, -0.2031],
                   [0.0000, 0.2031, 0.4062, 0.6250, 0.8125], [1.0000, 1.2500, 1.3750, 1.6250, 1.7500]]

    xq_scaled = [[-2.0000, -1.8571, -1.5714, -1.4286, -1.1429], [-1.0000, -0.7857, -0.5714, -0.3929, -0.1964],
                 [0.0000, 0.1964, 0.3929, 0.5714, 0.7857], [1.0000, 1.1429, 1.4286, 1.5714, 1.8571]]

    def test_e4m3_no_scale(self):
        x = torch.tensor(TestScaledE4M3.x, device="cuda")
        xq_ref = torch.tensor(TestScaledE4M3.xq_unscaled, device="cuda")
        e4m3_x = tensor_quant.scaled_e4m3(x, None)
        test_utils.compare(e4m3_x, xq_ref, atol=1e-4, rtol=1e-4)

    def test_e4m3_no_cpu(self):
        x = torch.tensor(TestScaledE4M3.x)
        xq_ref = torch.tensor(TestScaledE4M3.xq_unscaled)
        e4m3_x = tensor_quant.scaled_e4m3(x, None)
        test_utils.compare(e4m3_x, xq_ref, atol=1e-4, rtol=1e-4)

    def test_with_amax(self):
        x = torch.tensor(TestScaledE4M3.x, device="cuda").unsqueeze(-1)
        xq_ref = torch.tensor(TestScaledE4M3.xq_scaled, device="cuda").unsqueeze(-1)

        amax = quant_utils.reduce_amax(x, axis=None, keepdims=True)

        e4m3_x = tensor_quant.scaled_e4m3(x, amax)

        test_utils.compare(e4m3_x, xq_ref, atol=1e-4, rtol=1e-4)

    def test_e4m3_incontiguous(self):
        x = torch.tensor(TestScaledE4M3.x, device="cuda").transpose(1, 0)
        xq_ref = torch.tensor(TestScaledE4M3.xq_unscaled, device="cuda").transpose(1, 0)
        assert not x.is_contiguous()
        e4m3_x = tensor_quant.scaled_e4m3(x, None)
        test_utils.compare(e4m3_x, xq_ref, atol=1e-4, rtol=1e-4)

    def test_backward(self):
        x = torch.randn(3, 7, requires_grad=True).cuda()
        labels = torch.randint(6, (3,)).type(torch.LongTensor).cuda()
        quant_x = tensor_quant.scaled_e4m3(x, None)
        x.retain_grad()
        quant_x.retain_grad()
        criterion = torch.nn.CrossEntropyLoss().cuda()
        loss = criterion(quant_x, labels)
        loss.backward()
        np.testing.assert_array_equal(quant_x.grad.cpu().numpy(), x.grad.cpu().numpy())