brevitas.core package#

Subpackages#

Submodules#

brevitas.core.restrict_val module#

class brevitas.core.restrict_val.FloatRestrictValue[source]#

Bases: Module

combine_scale_threshold(x, threshold)[source]#
Return type:

Tensor

forward(x)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses. :rtype: Tensor

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.

restrict_init_float(x)[source]#
Return type:

float

restrict_init_inplace_module()[source]#
restrict_init_module()[source]#
restrict_init_tensor(x)[source]#
Return type:

Tensor

class brevitas.core.restrict_val.IntRestrictValue(restrict_value_float_to_int_impl=RoundSte())[source]#

Bases: Module

combine_scale_threshold(x, threshold)[source]#
Return type:

Tensor

forward(x)[source]#

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.

restrict_init_float(x)[source]#
restrict_init_inplace_module()[source]#
restrict_init_module()[source]#
restrict_init_tensor(x)[source]#
class brevitas.core.restrict_val.LogFloatRestrictValue[source]#

Bases: Module

combine_scale_threshold(x, threshold)[source]#
Return type:

Tensor

forward(x)[source]#

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.

restrict_init_float(x)[source]#
restrict_init_inplace_module()[source]#
restrict_init_module()[source]#
restrict_init_tensor(x)[source]#
class brevitas.core.restrict_val.PowerOfTwoRestrictValue(restrict_value_float_to_int_impl=RoundSte())[source]#

Bases: Module

combine_scale_threshold(x, threshold)[source]#
Return type:

Tensor

forward(x)[source]#

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.

restrict_init_float(x)[source]#
restrict_init_inplace_module()[source]#
restrict_init_module()[source]#
restrict_init_tensor(x)[source]#

brevitas.core.utils module#

class brevitas.core.utils.ParameterWrapper(value)[source]#

Bases: Module

forward()[source]#

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.

class brevitas.core.utils.SingleArgStatelessBuffer(value)[source]#

Bases: Module

forward(placeholder)[source]#

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.

class brevitas.core.utils.SliceTensor[source]#

Bases: Module

eager_forward(x)[source]#
Return type:

Tensor

forward(x)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses. :rtype: Tensor

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.

class brevitas.core.utils.StatelessBuffer(value)[source]#

Bases: Module

forward()[source]#

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.

state_dict(destination=None, prefix='', keep_vars=False)[source]#

Return a dictionary containing references to the whole state of the module.

Both parameters and persistent buffers (e.g. running averages) are included. Keys are corresponding parameter and buffer names. Parameters and buffers set to None are not included.

Note

The returned object is a shallow copy. It contains references to the module’s parameters and buffers.

Warning

Currently state_dict() also accepts positional arguments for destination, prefix and keep_vars in order. However, this is being deprecated and keyword arguments will be enforced in future releases.

Warning

Please avoid the use of argument destination as it is not designed for end-users.

Parameters:

  • destination (dict, optional) – If provided, the state of module will be updated into the dict and the same object is returned. Otherwise, an OrderedDict will be created and returned. Default: None.

  • prefix (str, optional) – a prefix added to parameter and buffer names to compose the keys in state_dict. Default: ''.

  • keep_vars (bool, optional) – by default the Tensor s returned in the state dict are detached from autograd. If it’s set to True, detaching will not be performed. Default: False.

Returns:

a dictionary containing a whole state of the module

Return type:

dict

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> module.state_dict().keys()
['bias', 'weight']
brevitas.core.utils.inplace_momentum_update(tensor, update, momentum, counter, new_counter)[source]#
Return type:

Tensor

brevitas.core.utils.inplace_tensor_add(tensor, value)[source]#
Return type:

Tensor

brevitas.core.utils.inplace_tensor_mul(tensor, value)[source]#
Return type:

Tensor

brevitas.core.zero_point module#

class brevitas.core.zero_point.ParameterFromRuntimeZeroPoint(collect_stats_steps, int_quant, quantize_zero_point, zero_point_stats_impl, zero_point_shape, zero_point_stats_input_view_shape_impl, zero_point_stats_momentum=0.1, dtype=None, device=None)[source]#

Bases: Module

forward(x, scale, bit_width)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses. :rtype: Tensor

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.

state_dict(destination=None, prefix='', keep_vars=False)[source]#

Return a dictionary containing references to the whole state of the module.

Both parameters and persistent buffers (e.g. running averages) are included. Keys are corresponding parameter and buffer names. Parameters and buffers set to None are not included.

Note

The returned object is a shallow copy. It contains references to the module’s parameters and buffers.

Warning

Currently state_dict() also accepts positional arguments for destination, prefix and keep_vars in order. However, this is being deprecated and keyword arguments will be enforced in future releases.

Warning

Please avoid the use of argument destination as it is not designed for end-users.

Parameters:

  • destination (dict, optional) – If provided, the state of module will be updated into the dict and the same object is returned. Otherwise, an OrderedDict will be created and returned. Default: None.

  • prefix (str, optional) – a prefix added to parameter and buffer names to compose the keys in state_dict. Default: ''.

  • keep_vars (bool, optional) – by default the Tensor s returned in the state dict are detached from autograd. If it’s set to True, detaching will not be performed. Default: False.

Returns:

a dictionary containing a whole state of the module

Return type:

dict

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> module.state_dict().keys()
['bias', 'weight']
training_forward(x)[source]#
Return type:

Tensor

class brevitas.core.zero_point.ParameterFromStatsFromParameterZeroPoint(int_quant, quantize_zero_point, zero_point_stats_input_view_shape_impl, zero_point_stats_input_concat_dim, zero_point_stats_impl, zero_point_shape, tracked_parameter_list, dtype=None, device=None)[source]#

Bases: Module

ScriptModule implementation of a learned scale factor initialized from statistics of a parameter, e.g. weights MSE or AbsMax.

forward(x, scale, bit_width)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses. :rtype: Tensor

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.

state_dict(destination=None, prefix='', keep_vars=False)[source]#

Return a dictionary containing references to the whole state of the module.

Both parameters and persistent buffers (e.g. running averages) are included. Keys are corresponding parameter and buffer names. Parameters and buffers set to None are not included.

Note

The returned object is a shallow copy. It contains references to the module’s parameters and buffers.

Warning

Currently state_dict() also accepts positional arguments for destination, prefix and keep_vars in order. However, this is being deprecated and keyword arguments will be enforced in future releases.

Warning

Please avoid the use of argument destination as it is not designed for end-users.

Parameters:

  • destination (dict, optional) – If provided, the state of module will be updated into the dict and the same object is returned. Otherwise, an OrderedDict will be created and returned. Default: None.

  • prefix (str, optional) – a prefix added to parameter and buffer names to compose the keys in state_dict. Default: ''.

  • keep_vars (bool, optional) – by default the Tensor s returned in the state dict are detached from autograd. If it’s set to True, detaching will not be performed. Default: False.

Returns:

a dictionary containing a whole state of the module

Return type:

dict

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> module.state_dict().keys()
['bias', 'weight']
class brevitas.core.zero_point.ParameterZeroPoint(zero_point_init, int_quant, quantize_zero_point, zero_point_shape=None, dtype=None, device=None)[source]#

Bases: Module

forward(x, scale, bit_width)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses. :rtype: Tensor

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.

class brevitas.core.zero_point.PreZeroCenterZeroPoint(stats_reduce_dim, pre_zero_point_stats_input_view_shape_impl, pre_zero_point_shape=None)[source]#

Bases: Module

Experimental ScriptModule implementation of a pre-scaling zero-point that zero-centers the incoming tensors. This is intended to be used with DecoupledIntQuant.

forward(x, scale, bit_width)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses. :rtype: Tensor

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.

get_zero_center(x)[source]#
Return type:

Tensor

class brevitas.core.zero_point.StatsFromParameterZeroPoint(int_quant, quantize_zero_point, zero_point_stats_input_view_shape_impl, zero_point_stats_input_concat_dim, zero_point_stats_impl, zero_point_shape, tracked_parameter_list)[source]#

Bases: Module

forward(x, scale, bit_width)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses. :rtype: Tensor

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.

class brevitas.core.zero_point.ZeroZeroPoint(dtype=None, device=None)[source]#

Bases: Module

forward(x, scale, bit_width)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses. :rtype: Tensor

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.

Module contents#