aihwkit.simulator.digital_low_precision.quantizers module

class aihwkit.simulator.digital_low_precision.quantizers.AsymmetricUniformQuantizer(n_bits, scale_domain='linear', per_channel=False, axis=None, eps=1e-08)

Bases: QuantizerBase

PyTorch Module that implements Asymmetric Uniform Quantization using STE. Quantizes its argument in the forward pass, passes the gradient ‘straight through’ on the backward pass, ignoring the quantization that occurred.

Parameters:

• n_bits (int) – Number of bits for quantization.

• scale_domain (str ('log', 'linear) with default='linear') – Domain of scale factor

• per_channel (bool) – If True: allows for per-channel quantization

property delta

forward(x_float)

Quantizes (quantized to integer and the scales back to original domain) :param x_float: Full-precision Tensor :type x_float: PyTorch Float Tensor

Returns:

x_quant – Quantized-Dequantized Tensor

Return type:

PyTorch Float Tensor

property int_max

property int_min

property is_initialized

make_range_trainable()

reset()

property scale

set_quant_range(x_min, x_max)

Instantiates the quantization parameters based on the provided min and max range

Parameters:

• x_min (tensor or float) – Quantization range minimum limit

• x_max (tensor of float) – Quantization range minimum limit

property symmetric

to_integer_forward(x_float)

Qunatized input to its integer represantion :param x_float: Full-precision Tensor :type x_float: PyTorch Float Tensor

Returns:

x_int

Return type:

PyTorch Float Tensor of integers

property x_max

property x_min

property zero_float

property zero_point

class aihwkit.simulator.digital_low_precision.quantizers.FloorStraightThrough(*args, **kwargs)

Bases: Function

static backward(ctx, output_grad)

Define a formula for differentiating the operation with backward mode automatic differentiation.

This function is to be overridden by all subclasses. (Defining this function is equivalent to defining the vjp function.)

It must accept a context ctx as the first argument, followed by as many outputs as the forward() returned (None will be passed in for non tensor outputs of the forward function), and it should return as many tensors, as there were inputs to forward(). Each argument is the gradient w.r.t the given output, and each returned value should be the gradient w.r.t. the corresponding input. If an input is not a Tensor or is a Tensor not requiring grads, you can just pass None as a gradient for that input.

The context can be used to retrieve tensors saved during the forward pass. It also has an attribute ctx.needs_input_grad as a tuple of booleans representing whether each input needs gradient. E.g., backward() will have ctx.needs_input_grad[0] = True if the first input to forward() needs gradient computed w.r.t. the output.

static forward(ctx, x)

Define the forward of the custom autograd Function.

This function is to be overridden by all subclasses. There are two ways to define forward:

Usage 1 (Combined forward and ctx):

@staticmethod
def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any:
    pass

• It must accept a context ctx as the first argument, followed by any number of arguments (tensors or other types).

• See combining-forward-context for more details

Usage 2 (Separate forward and ctx):

@staticmethod
def forward(*args: Any, **kwargs: Any) -> Any:
    pass


@staticmethod
def setup_context(ctx: Any, inputs: Tuple[Any, ...], output: Any) -> None:
    pass

• The forward no longer accepts a ctx argument.

• Instead, you must also override the torch.autograd.Function.setup_context() staticmethod to handle setting up the ctx object. output is the output of the forward, inputs are a Tuple of inputs to the forward.

• See extending-autograd for more details

The context can be used to store arbitrary data that can be then retrieved during the backward pass. Tensors should not be stored directly on ctx (though this is not currently enforced for backward compatibility). Instead, tensors should be saved either with ctx.save_for_backward() if they are intended to be used in backward (equivalently, vjp) or ctx.save_for_forward() if they are intended to be used for in jvp.

class aihwkit.simulator.digital_low_precision.quantizers.QMethodMap(value, cls)

Bases: tuple

cls

Alias for field number 1

value

Alias for field number 0

class aihwkit.simulator.digital_low_precision.quantizers.QMethods(*values)

Bases: Enum

asymmetric_uniform = (1, <class 'aihwkit.simulator.digital_low_precision.quantizers.AsymmetricUniformQuantizer'>)

property cls

classmethod list()

symmetric_uniform = (0, <class 'aihwkit.simulator.digital_low_precision.quantizers.SymmetricUniformQuantizer'>)

class aihwkit.simulator.digital_low_precision.quantizers.QuantizerBase(n_bits, per_channel=False, axis=None, *args, **kwargs)

Bases: Module

extra_repr()

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

forward(x_float)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

property is_initialized

reset()

set_quant_range(x_min, x_max)

property symmetric

property x_max

property x_min

exception aihwkit.simulator.digital_low_precision.quantizers.QuantizerNotInitializedError

Bases: Exception

Raised when a quantizer has not initialized

class aihwkit.simulator.digital_low_precision.quantizers.RoundStraightThrough(*args, **kwargs)

Bases: Function

static backward(ctx, output_grad)

Define a formula for differentiating the operation with backward mode automatic differentiation.

This function is to be overridden by all subclasses. (Defining this function is equivalent to defining the vjp function.)

It must accept a context ctx as the first argument, followed by as many outputs as the forward() returned (None will be passed in for non tensor outputs of the forward function), and it should return as many tensors, as there were inputs to forward(). Each argument is the gradient w.r.t the given output, and each returned value should be the gradient w.r.t. the corresponding input. If an input is not a Tensor or is a Tensor not requiring grads, you can just pass None as a gradient for that input.

The context can be used to retrieve tensors saved during the forward pass. It also has an attribute ctx.needs_input_grad as a tuple of booleans representing whether each input needs gradient. E.g., backward() will have ctx.needs_input_grad[0] = True if the first input to forward() needs gradient computed w.r.t. the output.

static forward(ctx, x)

Define the forward of the custom autograd Function.

This function is to be overridden by all subclasses. There are two ways to define forward:

Usage 1 (Combined forward and ctx):

@staticmethod
def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any:
    pass

• It must accept a context ctx as the first argument, followed by any number of arguments (tensors or other types).

• See combining-forward-context for more details

Usage 2 (Separate forward and ctx):

@staticmethod
def forward(*args: Any, **kwargs: Any) -> Any:
    pass


@staticmethod
def setup_context(ctx: Any, inputs: Tuple[Any, ...], output: Any) -> None:
    pass

• The forward no longer accepts a ctx argument.

• Instead, you must also override the torch.autograd.Function.setup_context() staticmethod to handle setting up the ctx object. output is the output of the forward, inputs are a Tuple of inputs to the forward.

• See extending-autograd for more details

The context can be used to store arbitrary data that can be then retrieved during the backward pass. Tensors should not be stored directly on ctx (though this is not currently enforced for backward compatibility). Instead, tensors should be saved either with ctx.save_for_backward() if they are intended to be used in backward (equivalently, vjp) or ctx.save_for_forward() if they are intended to be used for in jvp.

class aihwkit.simulator.digital_low_precision.quantizers.SymmetricUniformQuantizer(n_bits, scale_domain='linear', per_channel=False, axis=None, eps=1e-08)

Bases: QuantizerBase

PyTorch Module that implements Symmetric Uniform Quantization using STE. Quantizes its argument in the forward pass, passes the gradient ‘straight through’ on the backward pass, ignoring the quantization that occurred.

Parameters:

• n_bits (int) – Number of bits for quantization.

• scale_domain (str ('log', 'linear) with default='linear') – Domain of scale factor

• per_channel (bool) – If True: allows for per-channel quantization

property delta

forward(x_float)

Quantizes (quantized to integer and the scales back to original domain) :param x_float: Full-precision Tensor :type x_float: PyTorch Float Tensor

Returns:

x_quant – Quantized-Dequantized Tensor

Return type:

PyTorch Float Tensor

property int_max

property int_min

property is_initialized

make_range_trainable()

reset()

property scale

set_quant_range(x_min, x_max)

property signed

property symmetric

to_integer_forward(x_float)

Qunatized input to its integer represantion :param x_float: Full-precision Tensor :type x_float: PyTorch Float Tensor

Returns:

x_int

Return type:

PyTorch Float Tensor of integers

property x_max

property x_min

property zero_point