aihwkit.nn.low_precision_modules.quantized_input_module module

Module wrapper to quantize its input

class aihwkit.nn.low_precision_modules.quantized_input_module.QuantizedInputModule(module, act_quant_config)

Bases: Module

Wraps a module with an activation quantizer on the inputs, this offering the capability to quantize both the inputs and the outputs of a module. (NOTE: Output quantization is considered to be taken care of in the activation quantization of the module that is being wrapped)

This is useful when a module is a first layer and consumes directly from the dataloader as well as when a layer follows a functional operation (e.g., activation function or addition) which did not quantize its own output down to the required size.

Parameters:

• module (Module)

• act_quant_config (ActivationQuantConfig)

forward(inp, *args, **kwargs)

Perform the forward of the original module after quantizing its input

Parameters:

• inp (Tensor)

• args (Any)

• kwargs (Any)

Return type:

Tensor