from __future__ import annotations

from typing import Any

import paddle.fluid as fluid
from paddle.fluid import core as core

def fix_model_dict(model: Any): ...
def pre_hook(layer: Any, input: Any): ...
def post_hook(layer: Any, input: Any, output: Any): ...
def train_lenet(lenet: Any, reader: Any, optimizer: Any): ...

class ImperativeLenet(fluid.dygraph.Layer):
    features: Any = ...
    fc: Any = ...
    add: Any = ...
    quant_stub: Any = ...
    def __init__(self, num_classes: int = ...) -> None: ...
    def forward(self, inputs: Any): ...

class ImperativeLenetWithSkipQuant(fluid.dygraph.Layer):
    conv2d_0: Any = ...
    batch_norm_0: Any = ...
    relu_0: Any = ...
    pool2d_0: Any = ...
    conv2d_1: Any = ...
    batch_norm_1: Any = ...
    relu6_0: Any = ...
    pool2d_1: Any = ...
    linear_0: Any = ...
    leaky_relu_0: Any = ...
    linear_1: Any = ...
    sigmoid_0: Any = ...
    linear_2: Any = ...
    softmax_0: Any = ...
    def __init__(self, num_classes: int = ...) -> None: ...
    def forward(self, inputs: Any): ...

class ImperativeLinearBn(fluid.dygraph.Layer):
    linear: Any = ...
    bn: Any = ...
    def __init__(self) -> None: ...
    def forward(self, inputs: Any): ...

class ImperativeLinearBn_hook(fluid.dygraph.Layer):
    linear: Any = ...
    bn: Any = ...
    def __init__(self) -> None: ...
    def forward(self, inputs: Any): ...
