from __future__ import annotations

import unittest
from typing import Any

import paddle
import paddle.nn as nn
from paddle.fluid.contrib.slim.quantization import (
    QuantizationTransformPass as QuantizationTransformPass,
)
from paddle.nn.quant.quant_layers import (
    QuantizedConv2DTranspose as QuantizedConv2DTranspose,
)

class PACT(nn.Layer):
    alpha: Any = ...
    def __init__(self, init_value: int = ...) -> None: ...
    def forward(self, x: Any): ...

class CustomQAT(nn.Layer):
    u_param: Any = ...
    l_param: Any = ...
    alpha_param: Any = ...
    upper: Any = ...
    lower: Any = ...
    def __init__(self) -> None: ...
    def forward(self, x: Any): ...

class ModelForConv2dT(nn.Layer):
    features: Any = ...
    fc: Any = ...
    def __init__(self, num_classes: int = ...) -> None: ...
    def forward(self, inputs: Any): ...

class ImperativeLenet(paddle.nn.Layer):
    features: Any = ...
    fc: Any = ...
    def __init__(self, num_classes: int = ..., classifier_activation: str = ...) -> None: ...
    def forward(self, inputs: Any): ...

class TestUserDefinedActPreprocess(unittest.TestCase):
    imperative_qat: Any = ...
    def setUp(self) -> None: ...
    def func_quant_aware_training(self) -> None: ...
    def test_quant_aware_training(self) -> None: ...

class TestUserDefinedWeightPreprocess(TestUserDefinedActPreprocess):
    imperative_qat: Any = ...
    def setUp(self) -> None: ...

class TestUserDefinedActQuantize(TestUserDefinedActPreprocess):
    imperative_qat: Any = ...
    def setUp(self) -> None: ...

class TestUserDefinedWeightQuantize(TestUserDefinedActPreprocess):
    imperative_qat: Any = ...
    def setUp(self) -> None: ...
