From 1b9589dbe0bafc6f8a330b291912f8be5ccee68e Mon Sep 17 00:00:00 2001 From: Alessandro Pappalardo Date: Fri, 21 Apr 2023 17:31:32 +0100 Subject: [PATCH] Docs: update docs and READMEs Signed-off-by: Alessandro Pappalardo --- README.md | 3 +- docs/.buildinfo | 2 +- .../brevitas/core/bit_width/const.html | 101 +-- .../brevitas/core/bit_width/parameter.html | 92 +-- .../brevitas/core/function_wrapper/clamp.html | 47 +- .../brevitas/core/function_wrapper/misc.html | 52 +- .../core/function_wrapper/ops_ste.html | 46 +- .../brevitas/core/function_wrapper/shape.html | 49 +- docs/_modules/brevitas/core/quant/binary.html | 51 +- docs/_modules/brevitas/core/quant/delay.html | 48 +- docs/_modules/brevitas/core/quant/int.html | 146 ++-- .../brevitas/core/quant/int_base.html | 68 +- .../_modules/brevitas/core/quant/ternary.html | 47 +- docs/_modules/brevitas/core/restrict_val.html | 72 +- .../brevitas/core/scaling/int_scaling.html | 49 +- .../brevitas/core/scaling/runtime.html | 94 +-- .../brevitas/core/scaling/standalone.html | 103 +-- .../brevitas/core/stats/stats_op.html | 110 +-- docs/_modules/brevitas/core/utils.html | 40 +- docs/_modules/brevitas/core/zero_point.html | 135 ++-- docs/_modules/brevitas/function/ops.html | 47 +- docs/_modules/brevitas/function/ops_ste.html | 62 +- docs/_modules/brevitas/function/shape.html | 46 +- .../brevitas/ops/autograd_ste_ops.html | 68 +- docs/_modules/index.html | 2 +- docs/_static/documentation_options.js | 2 +- docs/about.html | 2 +- .../brevitas.core.bit_width.html | 103 ++- .../brevitas.core.function_wrapper.html | 54 +- docs/api_reference/brevitas.core.html | 360 ++++++++- docs/api_reference/brevitas.core.quant.html | 109 ++- docs/api_reference/brevitas.core.scaling.html | 36 +- docs/api_reference/brevitas.core.stats.html | 144 +++- docs/api_reference/brevitas.function.html | 50 +- docs/api_reference/brevitas.ops.html | 38 +- docs/api_reference/index.html | 2 +- docs/architecture.html | 2 +- docs/faq.html | 2 +- docs/genindex.html | 68 +- docs/getting_started.html | 81 +- docs/index.html | 2 +- docs/objects.inv | Bin 9964 -> 10518 bytes docs/py-modindex.html | 2 +- docs/search.html | 2 +- docs/searchindex.js | 2 +- docs/settings.html | 2 +- docs/setup.html | 2 +- docs/tutorials/anatomy_quantizer.html | 389 +++++----- docs/tutorials/anatomy_quantizer.ipynb | 402 +++++----- docs/tutorials/index.html | 2 +- docs/tutorials/onnx_export.html | 27 +- docs/tutorials/onnx_export.ipynb | 45 +- docs/tutorials/quant_activation_overview.html | 50 +- .../tutorials/quant_activation_overview.ipynb | 50 +- docs/tutorials/quant_recurrent.html | 26 +- docs/tutorials/quant_recurrent.ipynb | 24 +- .../quant_tensor_quant_conv2d_overview.html | 702 +++++++++--------- .../quant_tensor_quant_conv2d_overview.ipynb | 572 +++++++------- docs/tutorials/tvmcon2021.html | 6 +- docs/tutorials/tvmcon2021.ipynb | 4 +- docs/user_guide/datatypes.html | 2 +- docs/user_guide/export.html | 2 +- docs/user_guide/graph_transformations.html | 2 +- docs/user_guide/index.html | 2 +- docs/user_guide/precision.html | 2 +- docs/user_guide/proxies.html | 2 +- docs/user_guide/quantized_layers.html | 2 +- docs/user_guide/quantized_tensor.html | 2 +- docs/user_guide/quantizers.html | 2 +- docs/user_guide/scaling.html | 2 +- docs/user_guide/zero_point.html | 2 +- docsrc/source/getting_started.rst | 38 +- src/brevitas_examples/README.md | 7 + .../imagenet_classification/README.md | 16 +- .../imagenet_classification/models/README.md | 32 + .../imagenet_classification/ptq/README.md | 94 +-- .../imagenet_classification/qat/README.md | 20 - 77 files changed, 2585 insertions(+), 2588 deletions(-) create mode 100644 src/brevitas_examples/README.md create mode 100644 src/brevitas_examples/imagenet_classification/models/README.md diff --git a/README.md b/README.md index 6b37054df..003bba0f4 100644 --- a/README.md +++ b/README.md @@ -6,7 +6,7 @@ ![Examples Pytest](https://github.com/Xilinx/brevitas/workflows/Examples%20Pytest/badge.svg?branch=master) [![DOI](https://zenodo.org/badge/140494324.svg)](https://zenodo.org/badge/latestdoi/140494324) -Brevitas is a PyTorch library for neural network quantization, with a focus on *quantization-aware training (QAT)*. +Brevitas is a PyTorch library for neural network quantization, with support for both *post-training quantization (PTQ)* and *quantization-aware training (QAT)*. **Please note that Brevitas is a research project and not an official Xilinx product.** @@ -31,6 +31,7 @@ If you adopt Brevitas in your work, please cite it as: ## History +- *2023/04/21* - Release version 0.9.0, see the [release notes](https://github.com/Xilinx/brevitas/releases/tag/v0.9.0). - *2023/01/10* - Release version 0.8.0, see the [release notes](https://github.com/Xilinx/brevitas/releases/tag/v0.8.0). - *2021/12/13* - Release version 0.7.1, fix a bunch of issues. Added TVMCon 2021 tutorial notebook. - *2021/11/03* - Re-release version 0.7.0 (build 1) on PyPI to fix a packaging issue. diff --git a/docs/.buildinfo b/docs/.buildinfo index 920729fa6..9081611fe 100644 --- a/docs/.buildinfo +++ b/docs/.buildinfo @@ -1,4 +1,4 @@ # Sphinx build info version 1 # This file hashes the configuration used when building these files. When it is not found, a full rebuild will be done. -config: 50bc48ec3b3ea713612b0fb70b737f9c +config: bfc625bef8624d43d40fb6e0a790344b tags: 645f666f9bcd5a90fca523b33c5a78b7 diff --git a/docs/_modules/brevitas/core/bit_width/const.html b/docs/_modules/brevitas/core/bit_width/const.html index 2b4a20853..6739ff04b 100644 --- a/docs/_modules/brevitas/core/bit_width/const.html +++ b/docs/_modules/brevitas/core/bit_width/const.html @@ -5,7 +5,7 @@ - brevitas.core.bit_width.const — Brevitas 0.7.2.dev139+g0c2e90d documentation + brevitas.core.bit_width.const — Brevitas 0.9.0 documentation @@ -350,57 +350,21 @@

Source code for brevitas.core.bit_width.const

-# Copyright (c) 2018-     Xilinx, Inc              (Alessandro Pappalardo)
-# Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)
-# Copyright (c) 2014-     Facebook, Inc            (Soumith Chintala)
-# Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
-# Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)
-# Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
-# Copyright (c) 2011-2013 NYU                      (Clement Farabet)
-# Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)
-# Copyright (c) 2006      Idiap Research Institute (Samy Bengio)
-# Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
-
-# All rights reserved.
-
-# Redistribution and use in source and binary forms, with or without
-# modification, are permitted provided that the following conditions are met:
-
-# 1. Redistributions of source code must retain the above copyright
-#    notice, this list of conditions and the following disclaimer.
-
-# 2. Redistributions in binary form must reproduce the above copyright
-#    notice, this list of conditions and the following disclaimer in the
-#    documentation and/or other materials provided with the distribution.
-
-# 3. Neither the names of Xilinx, Facebook, Deepmind Technologies, NYU,
-#    NEC Laboratories America and IDIAP Research Institute nor the names
-#    of its contributors may be used to endorse or promote products derived
-#    from this software without specific prior written permission.
-
-# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
-# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
-# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
-# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
-# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
-# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
-# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
-# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
-# POSSIBILITY OF SUCH DAMAGE.
+# Copyright (C) 2023, Advanced Micro Devices, Inc. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
 
 import torch
 from torch import Tensor
 from torch.nn import Module
 
 import brevitas
-from brevitas.function.ops_ste import tensor_clamp_ste
+import brevitas.config as config
 from brevitas.core.utils import StatelessBuffer
+from brevitas.function.ops_ste import tensor_clamp_ste
 
 
 
[docs]class BitWidthConst(brevitas.jit.ScriptModule):
-    """ 
+    """
     ScriptModule that returns a constant bit-width wrapped in a float torch.tensor.
 
     Args:
@@ -418,6 +382,7 @@ 
Source code for brevitas.core.bit_width.const
    Note:
         Maps to bit_width_impl_type == BitWidthImplType.CONST == 'CONST' == 'const' in higher-level APIs.
     """
+
     def __init__(self, bit_width: int) -> None:
         super(BitWidthConst, self).__init__()
         assert isinstance(bit_width, int)
@@ -428,6 +393,54 @@ 
Source code for brevitas.core.bit_width.const
return self.bit_width()

 
 
+
[docs]class BitWidthStatefulConst(brevitas.jit.ScriptModule):
+    """
+    ScriptModule that returns a constant bit-width wrapped in a float torch.tensor but retains the
+    bit-width as part of the module state.
+
+    Args:
+        bit_width (int): bit-width value.
+
+    Examples:
+        >>> bit_width = BitWidthStatefulConst(8)
+        >>> bit_width()
+        tensor(8.)
+
+    Note:
+        The BitWidthStatefulConst is a counterpart to BitWidthConst with the difference that the
+        BitWidthStatefulConst retains the bit-width as part of the Module's state. This means that it
+        will be saved as part of a checkpoint.
+
+    Note:
+        Maps to bit_width_impl_type == BitWidthImplType.STATEFUL_CONST == 'STATEFUL_CONST' ==
+        'stateful_const' in higher-level APIs.
+    """
+
+    def __init__(self, bit_width: int) -> None:
+        super(BitWidthStatefulConst, self).__init__()
+        assert isinstance(bit_width, int)
+        self.register_buffer("bit_width", torch.tensor(float(bit_width)))
+
+
[docs]    @brevitas.jit.script_method
+    def forward(self) -> Tensor:
+        return self.bit_width

+
+    def _load_from_state_dict(
+            self,
+            state_dict,
+            prefix,
+            local_metadata,
+            strict,
+            missing_keys: list,
+            unexpected_keys,
+            error_msgs):
+        super(BitWidthStatefulConst, self)._load_from_state_dict(
+            state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
+        value_key = prefix + "bit_width"
+        if config.IGNORE_MISSING_KEYS and value_key in missing_keys:
+            missing_keys.remove(value_key)

+
+
 
[docs]class MsbClampBitWidth(brevitas.jit.ScriptModule):
 
     def __init__(
@@ -446,9 +459,7 @@ 
Source code for brevitas.core.bit_width.const
bit_width_to_remove = self.bit_width_to_remove_impl()
         output_bit_width = torch.abs(input_bit_width - bit_width_to_remove)
         output_bit_width = tensor_clamp_ste(
-            output_bit_width,
-            self.min_overall_bit_width(),
-            self.max_overall_bit_width())
+            output_bit_width, self.min_overall_bit_width(), self.max_overall_bit_width())
         return output_bit_width
diff --git a/docs/_modules/brevitas/core/bit_width/parameter.html b/docs/_modules/brevitas/core/bit_width/parameter.html index 0c9e7553e..a467dc162 100644 --- a/docs/_modules/brevitas/core/bit_width/parameter.html +++ b/docs/_modules/brevitas/core/bit_width/parameter.html @@ -5,7 +5,7 @@ - brevitas.core.bit_width.parameter — Brevitas 0.7.2.dev139+g0c2e90d documentation + brevitas.core.bit_width.parameter — Brevitas 0.9.0 documentation @@ -350,58 +350,19 @@

Source code for brevitas.core.bit_width.parameter

-# Copyright (c) 2018-     Xilinx, Inc              (Alessandro Pappalardo)
-# Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)
-# Copyright (c) 2014-     Facebook, Inc            (Soumith Chintala)
-# Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
-# Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)
-# Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
-# Copyright (c) 2011-2013 NYU                      (Clement Farabet)
-# Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)
-# Copyright (c) 2006      Idiap Research Institute (Samy Bengio)
-# Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
-
-# All rights reserved.
-
-# Redistribution and use in source and binary forms, with or without
-# modification, are permitted provided that the following conditions are met:
-
-# 1. Redistributions of source code must retain the above copyright
-#    notice, this list of conditions and the following disclaimer.
-
-# 2. Redistributions in binary form must reproduce the above copyright
-#    notice, this list of conditions and the following disclaimer in the
-#    documentation and/or other materials provided with the distribution.
-
-# 3. Neither the names of Xilinx, Facebook, Deepmind Technologies, NYU,
-#    NEC Laboratories America and IDIAP Research Institute nor the names
-#    of its contributors may be used to endorse or promote products derived
-#    from this software without specific prior written permission.
-
-# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
-# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
-# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
-# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
-# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
-# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
-# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
-# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
-# POSSIBILITY OF SUCH DAMAGE.
-
-from typing import Optional
+# Copyright (C) 2023, Advanced Micro Devices, Inc. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
 
 import torch
 from torch import Tensor
-from torch.nn import Parameter, Module
+from torch.nn import Module
+from torch.nn import Parameter
 
 import brevitas
 import brevitas.config as config
-from brevitas.function import abs_binary_sign_grad
 from brevitas.core.function_wrapper import RoundSte
 from brevitas.core.restrict_val import IntRestrictValue
-
+from brevitas.function import abs_binary_sign_grad
 
 MIN_INT_BIT_WIDTH = 2
 NON_ZERO_EPSILON = 1e-6
@@ -447,16 +408,19 @@ 
Source code for brevitas.core.bit_width.parameter
super(BitWidthParameter, self).__init__()
 
         if bit_width < MIN_INT_BIT_WIDTH:
-            raise RuntimeError("Int bit width has to be at least {}, instead is {}."
-                            .format(MIN_INT_BIT_WIDTH, bit_width))
+            raise RuntimeError(
+                "Int bit width has to be at least {}, instead is {}.".format(
+                    MIN_INT_BIT_WIDTH, bit_width))
 
         if min_bit_width < MIN_INT_BIT_WIDTH:
-            raise RuntimeError("Min int bit width has to be at least {}, instead is {}."
-                            .format(MIN_INT_BIT_WIDTH, min_bit_width))
+            raise RuntimeError(
+                "Min int bit width has to be at least {}, instead is {}.".format(
+                    MIN_INT_BIT_WIDTH, min_bit_width))
 
         if bit_width < min_bit_width:
-            raise RuntimeError("Int bit width has to be at least {}, instead is {}."
-                            .format(min_bit_width, bit_width))
+            raise RuntimeError(
+                "Int bit width has to be at least {}, instead is {}.".format(
+                    min_bit_width, bit_width))
 
         bit_width = float(int(bit_width))
         min_bit_width = float(int(min_bit_width))
@@ -474,13 +438,21 @@ 
Source code for brevitas.core.bit_width.parameter
bit_width = self.restrict_bit_width_impl(bit_width)
         return bit_width

 
-    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict,
-                              missing_keys, unexpected_keys, error_msgs):
+    def _load_from_state_dict(
+            self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys,
+            error_msgs):
+        bit_width_const_key = prefix + 'bit_width'
         bit_width_offset_key = prefix + 'bit_width_offset'
+        if bit_width_const_key in state_dict:
+            assert bit_width_offset_key not in state_dict, "Both should not be true."
+            bit_width = state_dict[bit_width_const_key]
+            state_dict[bit_width_offset_key] = bit_width - self.bit_width_base
+            del state_dict[bit_width_const_key]
+
         if self.override_pretrained and bit_width_offset_key in state_dict:
             del state_dict[bit_width_offset_key]
         super(BitWidthParameter, self)._load_from_state_dict(
-            state_dict, prefix, local_metadata,strict,missing_keys, unexpected_keys, error_msgs)
+            state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
         if config.IGNORE_MISSING_KEYS and bit_width_offset_key in missing_keys:
             missing_keys.remove(bit_width_offset_key)

 
@@ -493,7 +465,7 @@ 
Source code for brevitas.core.bit_width.parameter
bit_width_to_remove: int,
             override_pretrained_bit_width: bool = False,
             non_zero_epsilon: float = NON_ZERO_EPSILON,
-            remove_zero_bit_width = REMOVE_ZERO_BIT_WIDTH):
+            remove_zero_bit_width=REMOVE_ZERO_BIT_WIDTH):
         super(RemoveBitwidthParameter, self).__init__()
 
         if bit_width_to_remove < 0:
@@ -511,8 +483,9 @@ 
Source code for brevitas.core.bit_width.parameter
bit_width_to_remove = 1.0 / (self.non_zero_epsilon + torch.abs(self.bit_width_coeff))
         return bit_width_to_remove

 
-    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict,
-                              missing_keys, unexpected_keys, error_msgs):
+    def _load_from_state_dict(
+            self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys,
+            error_msgs):
         bit_width_coeff_key = prefix + 'bit_width_coeff'
         if self.override_pretrained and bit_width_coeff_key in state_dict:
             del state_dict[bit_width_coeff_key]
@@ -520,11 +493,6 @@ 
Source code for brevitas.core.bit_width.parameter
state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
         if config.IGNORE_MISSING_KEYS and bit_width_coeff_key in missing_keys:
             missing_keys.remove(bit_width_coeff_key)

-
-
-
-
-
diff --git a/docs/_modules/brevitas/core/function_wrapper/clamp.html b/docs/_modules/brevitas/core/function_wrapper/clamp.html index 94fdfec51..5b881504a 100644 --- a/docs/_modules/brevitas/core/function_wrapper/clamp.html +++ b/docs/_modules/brevitas/core/function_wrapper/clamp.html @@ -5,7 +5,7 @@ - brevitas.core.function_wrapper.clamp — Brevitas 0.7.2.dev139+g0c2e90d documentation + brevitas.core.function_wrapper.clamp — Brevitas 0.9.0 documentation @@ -350,54 +350,17 @@

Source code for brevitas.core.function_wrapper.clamp

-# Copyright (c) 2018-     Xilinx, Inc              (Alessandro Pappalardo)
-# Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)
-# Copyright (c) 2014-     Facebook, Inc            (Soumith Chintala)
-# Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
-# Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)
-# Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
-# Copyright (c) 2011-2013 NYU                      (Clement Farabet)
-# Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)
-# Copyright (c) 2006      Idiap Research Institute (Samy Bengio)
-# Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
-
-# All rights reserved.
-
-# Redistribution and use in source and binary forms, with or without
-# modification, are permitted provided that the following conditions are met:
-
-# 1. Redistributions of source code must retain the above copyright
-#    notice, this list of conditions and the following disclaimer.
-
-# 2. Redistributions in binary form must reproduce the above copyright
-#    notice, this list of conditions and the following disclaimer in the
-#    documentation and/or other materials provided with the distribution.
-
-# 3. Neither the names of Xilinx, Facebook, Deepmind Technologies, NYU,
-#    NEC Laboratories America and IDIAP Research Institute nor the names
-#    of its contributors may be used to endorse or promote products derived
-#    from this software without specific prior written permission.
-
-# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
-# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
-# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
-# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
-# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
-# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
-# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
-# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
-# POSSIBILITY OF SUCH DAMAGE.
+# Copyright (C) 2023, Advanced Micro Devices, Inc. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
 
 """
 ScriptModule wrappers for various variants of clamping.
 """
 
-
-import brevitas
 import torch
 from torch import Tensor
+
+import brevitas
 from brevitas.function import tensor_clamp
 
 
diff --git a/docs/_modules/brevitas/core/function_wrapper/misc.html b/docs/_modules/brevitas/core/function_wrapper/misc.html
index ab4a257d1..5295fdcc8 100644
--- a/docs/_modules/brevitas/core/function_wrapper/misc.html
+++ b/docs/_modules/brevitas/core/function_wrapper/misc.html
@@ -5,7 +5,7 @@
   
     
     
-    brevitas.core.function_wrapper.misc — Brevitas 0.7.2.dev139+g0c2e90d documentation
+    brevitas.core.function_wrapper.misc — Brevitas 0.9.0 documentation
   
   
   
@@ -350,53 +350,17 @@
             

               
   
Source code for brevitas.core.function_wrapper.misc
-# Copyright (c) 2018-     Xilinx, Inc              (Alessandro Pappalardo)
-# Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)
-# Copyright (c) 2014-     Facebook, Inc            (Soumith Chintala)
-# Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
-# Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)
-# Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
-# Copyright (c) 2011-2013 NYU                      (Clement Farabet)
-# Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)
-# Copyright (c) 2006      Idiap Research Institute (Samy Bengio)
-# Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
-
-# All rights reserved.
-
-# Redistribution and use in source and binary forms, with or without
-# modification, are permitted provided that the following conditions are met:
-
-# 1. Redistributions of source code must retain the above copyright
-#    notice, this list of conditions and the following disclaimer.
-
-# 2. Redistributions in binary form must reproduce the above copyright
-#    notice, this list of conditions and the following disclaimer in the
-#    documentation and/or other materials provided with the distribution.
-
-# 3. Neither the names of Xilinx, Facebook, Deepmind Technologies, NYU,
-#    NEC Laboratories America and IDIAP Research Institute nor the names
-#    of its contributors may be used to endorse or promote products derived
-#    from this software without specific prior written permission.
-
-# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
-# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
-# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
-# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
-# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
-# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
-# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
-# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
-# POSSIBILITY OF SUCH DAMAGE.
+# Copyright (C) 2023, Advanced Micro Devices, Inc. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
 
 """
 A collection of miscellaneous ScriptModule used in various quantizers.
 """
 
-import brevitas
 import torch
 
+import brevitas
+
 
 
[docs]class Identity(brevitas.jit.ScriptModule):
     """
@@ -478,12 +442,6 @@ 
Source code for brevitas.core.function_wrapper.misc
def forward(self, x: torch.Tensor) -> torch.Tensor:
         x.log2_()
         return x

-
-
-
-
-
-
 

 
             

diff --git a/docs/_modules/brevitas/core/function_wrapper/ops_ste.html b/docs/_modules/brevitas/core/function_wrapper/ops_ste.html
index fa215728f..45d6c3e1b 100644
--- a/docs/_modules/brevitas/core/function_wrapper/ops_ste.html
+++ b/docs/_modules/brevitas/core/function_wrapper/ops_ste.html
@@ -5,7 +5,7 @@
   
     
     
-    brevitas.core.function_wrapper.ops_ste — Brevitas 0.7.2.dev139+g0c2e90d documentation
+    brevitas.core.function_wrapper.ops_ste — Brevitas 0.9.0 documentation
   
   
   
@@ -350,52 +350,16 @@
             

               
   
Source code for brevitas.core.function_wrapper.ops_ste
-# Copyright (c) 2018-     Xilinx, Inc              (Alessandro Pappalardo)
-# Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)
-# Copyright (c) 2014-     Facebook, Inc            (Soumith Chintala)
-# Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
-# Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)
-# Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
-# Copyright (c) 2011-2013 NYU                      (Clement Farabet)
-# Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)
-# Copyright (c) 2006      Idiap Research Institute (Samy Bengio)
-# Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
-
-# All rights reserved.
-
-# Redistribution and use in source and binary forms, with or without
-# modification, are permitted provided that the following conditions are met:
-
-# 1. Redistributions of source code must retain the above copyright
-#    notice, this list of conditions and the following disclaimer.
-
-# 2. Redistributions in binary form must reproduce the above copyright
-#    notice, this list of conditions and the following disclaimer in the
-#    documentation and/or other materials provided with the distribution.
-
-# 3. Neither the names of Xilinx, Facebook, Deepmind Technologies, NYU,
-#    NEC Laboratories America and IDIAP Research Institute nor the names
-#    of its contributors may be used to endorse or promote products derived
-#    from this software without specific prior written permission.
-
-# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
-# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
-# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
-# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
-# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
-# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
-# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
-# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
-# POSSIBILITY OF SUCH DAMAGE.
+# Copyright (C) 2023, Advanced Micro Devices, Inc. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
 
 """
 ScriptModule wrappers of various functions defined in :obj:`~brevitas.function.ops_ste`.
 """
 
-import brevitas
 import torch
+
+import brevitas
 from brevitas.function.ops_ste import *
 
 
diff --git a/docs/_modules/brevitas/core/function_wrapper/shape.html b/docs/_modules/brevitas/core/function_wrapper/shape.html
index ae8f69891..9fed2caf1 100644
--- a/docs/_modules/brevitas/core/function_wrapper/shape.html
+++ b/docs/_modules/brevitas/core/function_wrapper/shape.html
@@ -5,7 +5,7 @@
   
     
     
-    brevitas.core.function_wrapper.shape — Brevitas 0.7.2.dev139+g0c2e90d documentation
+    brevitas.core.function_wrapper.shape — Brevitas 0.9.0 documentation
   
   
   
@@ -350,45 +350,8 @@
             

               
   
Source code for brevitas.core.function_wrapper.shape
-# Copyright (c) 2018-     Xilinx, Inc              (Alessandro Pappalardo)
-# Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)
-# Copyright (c) 2014-     Facebook, Inc            (Soumith Chintala)
-# Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
-# Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)
-# Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
-# Copyright (c) 2011-2013 NYU                      (Clement Farabet)
-# Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)
-# Copyright (c) 2006      Idiap Research Institute (Samy Bengio)
-# Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
-
-# All rights reserved.
-
-# Redistribution and use in source and binary forms, with or without
-# modification, are permitted provided that the following conditions are met:
-
-# 1. Redistributions of source code must retain the above copyright
-#    notice, this list of conditions and the following disclaimer.
-
-# 2. Redistributions in binary form must reproduce the above copyright
-#    notice, this list of conditions and the following disclaimer in the
-#    documentation and/or other materials provided with the distribution.
-
-# 3. Neither the names of Xilinx, Facebook, Deepmind Technologies, NYU,
-#    NEC Laboratories America and IDIAP Research Institute nor the names
-#    of its contributors may be used to endorse or promote products derived
-#    from this software without specific prior written permission.
-
-# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
-# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
-# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
-# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
-# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
-# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
-# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
-# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
-# POSSIBILITY OF SUCH DAMAGE.
+# Copyright (C) 2023, Advanced Micro Devices, Inc. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
 
 """
 ScriptModule classes to compute the view of a tensor according to various different criteria.
@@ -399,9 +362,11 @@ 
Source code for brevitas.core.function_wrapper.shape
import torch
 
 import brevitas
-from brevitas.function.shape import over_tensor, over_output_channels, over_batch_over_tensor
-from brevitas.function.shape import over_batch_over_output_channels
 from brevitas.core.function_wrapper import Identity
+from brevitas.function.shape import over_batch_over_output_channels
+from brevitas.function.shape import over_batch_over_tensor
+from brevitas.function.shape import over_output_channels
+from brevitas.function.shape import over_tensor
 
 
 
[docs]class PermuteDims(brevitas.jit.ScriptModule):
diff --git a/docs/_modules/brevitas/core/quant/binary.html b/docs/_modules/brevitas/core/quant/binary.html
index 7a3c20525..d37cb0573 100644
--- a/docs/_modules/brevitas/core/quant/binary.html
+++ b/docs/_modules/brevitas/core/quant/binary.html
@@ -5,7 +5,7 @@
   
     
     
-    brevitas.core.quant.binary — Brevitas 0.7.2.dev139+g0c2e90d documentation
+    brevitas.core.quant.binary — Brevitas 0.9.0 documentation
   
   
   
@@ -350,45 +350,8 @@
             

               
   
Source code for brevitas.core.quant.binary
-# Copyright (c) 2018-     Xilinx, Inc              (Alessandro Pappalardo)
-# Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)
-# Copyright (c) 2014-     Facebook, Inc            (Soumith Chintala)
-# Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
-# Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)
-# Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
-# Copyright (c) 2011-2013 NYU                      (Clement Farabet)
-# Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)
-# Copyright (c) 2006      Idiap Research Institute (Samy Bengio)
-# Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
-
-# All rights reserved.
-
-# Redistribution and use in source and binary forms, with or without
-# modification, are permitted provided that the following conditions are met:
-
-# 1. Redistributions of source code must retain the above copyright
-#    notice, this list of conditions and the following disclaimer.
-
-# 2. Redistributions in binary form must reproduce the above copyright
-#    notice, this list of conditions and the following disclaimer in the
-#    documentation and/or other materials provided with the distribution.
-
-# 3. Neither the names of Xilinx, Facebook, Deepmind Technologies, NYU,
-#    NEC Laboratories America and IDIAP Research Institute nor the names
-#    of its contributors may be used to endorse or promote products derived
-#    from this software without specific prior written permission.
-
-# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
-# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
-# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
-# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
-# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
-# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
-# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
-# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
-# POSSIBILITY OF SUCH DAMAGE.
+# Copyright (C) 2023, Advanced Micro Devices, Inc. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
 
 from typing import Tuple
 
@@ -397,11 +360,11 @@ 
Source code for brevitas.core.quant.binary
 from torch.nn import Module
 
 import brevitas
-from brevitas.function.ops_ste import binary_sign_ste
-from brevitas.core.function_wrapper import TensorClamp
 from brevitas.core.bit_width import BitWidthConst
-from brevitas.core.utils import StatelessBuffer
+from brevitas.core.function_wrapper import TensorClamp
 from brevitas.core.quant.delay import DelayWrapper
+from brevitas.core.utils import StatelessBuffer
+from brevitas.function.ops_ste import binary_sign_ste
 
 
 
[docs]class BinaryQuant(brevitas.jit.ScriptModule):
@@ -507,7 +470,7 @@ 
Source code for brevitas.core.quant.binary
 
[docs]    @brevitas.jit.script_method
     def forward(self, x: Tensor) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
         scale = self.scaling_impl(x)
-        y = self.tensor_clamp_impl(x, - scale, scale)
+        y = self.tensor_clamp_impl(x, -scale, scale)
         y = binary_sign_ste(y) * scale
         y = self.delay_wrapper(x, y)
         return y, scale, self.zero_point(), self.bit_width()

diff --git a/docs/_modules/brevitas/core/quant/delay.html b/docs/_modules/brevitas/core/quant/delay.html
index 413a7cf9e..0c71380e9 100644
--- a/docs/_modules/brevitas/core/quant/delay.html
+++ b/docs/_modules/brevitas/core/quant/delay.html
@@ -5,7 +5,7 @@
   
     
     
-    brevitas.core.quant.delay — Brevitas 0.7.2.dev139+g0c2e90d documentation
+    brevitas.core.quant.delay — Brevitas 0.9.0 documentation
   
   
   
@@ -350,45 +350,8 @@
             

               
   
Source code for brevitas.core.quant.delay
-# Copyright (c) 2018-     Xilinx, Inc              (Alessandro Pappalardo)
-# Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)
-# Copyright (c) 2014-     Facebook, Inc            (Soumith Chintala)
-# Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
-# Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)
-# Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
-# Copyright (c) 2011-2013 NYU                      (Clement Farabet)
-# Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)
-# Copyright (c) 2006      Idiap Research Institute (Samy Bengio)
-# Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
-
-# All rights reserved.
-
-# Redistribution and use in source and binary forms, with or without
-# modification, are permitted provided that the following conditions are met:
-
-# 1. Redistributions of source code must retain the above copyright
-#    notice, this list of conditions and the following disclaimer.
-
-# 2. Redistributions in binary form must reproduce the above copyright
-#    notice, this list of conditions and the following disclaimer in the
-#    documentation and/or other materials provided with the distribution.
-
-# 3. Neither the names of Xilinx, Facebook, Deepmind Technologies, NYU,
-#    NEC Laboratories America and IDIAP Research Institute nor the names
-#    of its contributors may be used to endorse or promote products derived
-#    from this software without specific prior written permission.
-
-# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
-# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
-# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
-# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
-# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
-# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
-# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
-# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
-# POSSIBILITY OF SUCH DAMAGE.
+# Copyright (C) 2023, Advanced Micro Devices, Inc. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
 
 from typing import Optional
 
@@ -418,8 +381,9 @@ 
Source code for brevitas.core.quant.delay
         else:
             return y
 
-    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict,
-                              missing_keys, unexpected_keys, error_msgs):
+    def _load_from_state_dict(
+            self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys,
+            error_msgs):
         super(_DelayQuant, self)._load_from_state_dict(
             state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
         # Pytorch stores training flag as a buffer with JIT enabled
diff --git a/docs/_modules/brevitas/core/quant/int.html b/docs/_modules/brevitas/core/quant/int.html
index 8fa4aea39..ced9d191e 100644
--- a/docs/_modules/brevitas/core/quant/int.html
+++ b/docs/_modules/brevitas/core/quant/int.html
@@ -5,7 +5,7 @@
   
     
     
-    brevitas.core.quant.int — Brevitas 0.7.2.dev139+g0c2e90d documentation
+    brevitas.core.quant.int — Brevitas 0.9.0 documentation
   
   
   
@@ -350,45 +350,8 @@
             

               
   
Source code for brevitas.core.quant.int
-# Copyright (c) 2018-     Xilinx, Inc              (Alessandro Pappalardo)
-# Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)
-# Copyright (c) 2014-     Facebook, Inc            (Soumith Chintala)
-# Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
-# Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)
-# Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
-# Copyright (c) 2011-2013 NYU                      (Clement Farabet)
-# Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)
-# Copyright (c) 2006      Idiap Research Institute (Samy Bengio)
-# Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
-
-# All rights reserved.
-
-# Redistribution and use in source and binary forms, with or without
-# modification, are permitted provided that the following conditions are met:
-
-# 1. Redistributions of source code must retain the above copyright
-#    notice, this list of conditions and the following disclaimer.
-
-# 2. Redistributions in binary form must reproduce the above copyright
-#    notice, this list of conditions and the following disclaimer in the
-#    documentation and/or other materials provided with the distribution.
-
-# 3. Neither the names of Xilinx, Facebook, Deepmind Technologies, NYU,
-#    NEC Laboratories America and IDIAP Research Institute nor the names
-#    of its contributors may be used to endorse or promote products derived
-#    from this software without specific prior written permission.
-
-# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
-# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
-# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
-# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
-# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
-# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
-# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
-# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
-# POSSIBILITY OF SUCH DAMAGE.
+# Copyright (C) 2023, Advanced Micro Devices, Inc. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
 
 from typing import Optional, Tuple
 
@@ -397,9 +360,9 @@ 
Source code for brevitas.core.quant.int
 from torch.nn import Module
 
 import brevitas
-from brevitas.function.ops_ste import round_ste
-from brevitas.core.utils import StatelessBuffer
 from brevitas.core.quant.delay import DelayWrapper
+from brevitas.core.utils import StatelessBuffer
+from brevitas.function.ops_ste import round_ste
 
 
 
[docs]class PrescaledRestrictIntQuantWithInputBitWidth(brevitas.jit.ScriptModule):
@@ -438,18 +401,15 @@ 
Source code for brevitas.core.quant.int
     Note:
         Set env variable BREVITAS_JIT=1 to enable TorchScript compilation of this module.
     """
-    def __init__(self,
-                 int_quant: Module,
-                 bit_width_impl: Module):
+
+    def __init__(self, int_quant: Module, bit_width_impl: Module):
         super(PrescaledRestrictIntQuantWithInputBitWidth, self).__init__()
         self.int_quant = int_quant
         self.msb_clamp_bit_width_impl = bit_width_impl
         self.zero_point = StatelessBuffer(torch.tensor(0.0))
 
 
[docs]    @brevitas.jit.script_method
-    def forward(self,
-                x: Tensor,
-                scale: Tensor,
+    def forward(self, x: Tensor, scale: Tensor,
                 input_bit_width: Tensor) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
         bit_width = self.msb_clamp_bit_width_impl(input_bit_width)
         zero_point = self.zero_point()
@@ -460,19 +420,15 @@ 
Source code for brevitas.core.quant.int
 
[docs]class PrescaledRestrictIntQuant(brevitas.jit.ScriptModule):
     """
     """
-    def __init__(self,
-                 int_quant: Module,
-                 bit_width_impl: Module):
+
+    def __init__(self, int_quant: Module, bit_width_impl: Module):
         super(PrescaledRestrictIntQuant, self).__init__()
         self.int_quant = int_quant
         self.msb_clamp_bit_width_impl = bit_width_impl
         self.zero_point = StatelessBuffer(torch.tensor(0.0))
 
-
 
[docs]    @brevitas.jit.script_method
-    def forward(self,
-                x: Tensor,
-                scale: Tensor) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
+    def forward(self, x: Tensor, scale: Tensor) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
         msb_clamp_bit_width = self.msb_clamp_bit_width_impl()
         zero_point = self.zero_point()
         y = self.int_quant(scale, zero_point, msb_clamp_bit_width, x)
@@ -527,12 +483,14 @@ 
Source code for brevitas.core.quant.int
     Note:
         Set env variable BREVITAS_JIT=1 to enable TorchScript compilation of this module.
     """
-    def __init__(self,
-                 int_quant: Module,
-                 scaling_impl: Module,
-                 int_scaling_impl: Module,
-                 zero_point_impl: Module,
-                 bit_width_impl: Module):
+
+    def __init__(
+            self,
+            int_quant: Module,
+            scaling_impl: Module,
+            int_scaling_impl: Module,
+            zero_point_impl: Module,
+            bit_width_impl: Module):
         super(RescalingIntQuant, self).__init__()
         self.int_quant = int_quant
         self.scaling_impl = scaling_impl
@@ -553,14 +511,15 @@ 
Source code for brevitas.core.quant.int
 
 
[docs]class DecoupledRescalingIntQuant(brevitas.jit.ScriptModule):
 
-    def __init__(self,
-                 decoupled_int_quant: Module,
-                 pre_scaling_impl: Module,
-                 scaling_impl: Module,
-                 int_scaling_impl: Module,
-                 pre_zero_point_impl: Module,
-                 zero_point_impl: Module,
-                 bit_width_impl: Module):
+    def __init__(
+            self,
+            decoupled_int_quant: Module,
+            pre_scaling_impl: Module,
+            scaling_impl: Module,
+            int_scaling_impl: Module,
+            pre_zero_point_impl: Module,
+            zero_point_impl: Module,
+            bit_width_impl: Module):
         super(DecoupledRescalingIntQuant, self).__init__()
         self.decoupled_int_quant = decoupled_int_quant
         self.pre_scaling_impl = pre_scaling_impl
@@ -587,21 +546,16 @@ 
Source code for brevitas.core.quant.int
 
[docs]class TruncIntQuant(brevitas.jit.ScriptModule):
     """
     """
+
     def __init__(
-            self,
-            float_to_int_impl: Module,
-            bit_width_impl: Module,
-            quant_delay_steps: int = 0):
+            self, float_to_int_impl: Module, bit_width_impl: Module, quant_delay_steps: int = 0):
         super(TruncIntQuant, self).__init__()
         self.msb_clamp_bit_width_impl = bit_width_impl
         self.float_to_int_impl = float_to_int_impl
         self.delay_wrapper = DelayWrapper(quant_delay_steps)
 
 
[docs]    @brevitas.jit.script_method
-    def forward(self,
-                x: Tensor,
-                scale: Tensor,
-                zero_point: Tensor,
+    def forward(self, x: Tensor, scale: Tensor, zero_point: Tensor,
                 input_bit_width: Tensor) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
         y = x / scale
         y = y + zero_point
@@ -615,6 +569,44 @@ 
Source code for brevitas.core.quant.int
         y = y * scale
         y = self.delay_wrapper(x, y)
         return y, scale, zero_point, output_bit_width

+
+
+
[docs]class DecoupledRescalingIntQuantWithInput(DecoupledRescalingIntQuant):
+
+    def __init__(
+        self,
+        decoupled_int_quant: Module,
+        pre_scaling_impl: Module,
+        scaling_impl: Module,
+        int_scaling_impl: Module,
+        pre_zero_point_impl: Module,
+        zero_point_impl: Module,
+        bit_width_impl: Module,
+    ):
+        super().__init__(
+            decoupled_int_quant,
+            pre_scaling_impl,
+            scaling_impl,
+            int_scaling_impl,
+            pre_zero_point_impl,
+            zero_point_impl,
+            bit_width_impl,
+        )
+        # TODO - check the make sure the pre-scaling module takes the input bit-width and sign
+
+
[docs]    @brevitas.jit.script_method
+    def forward(self, x: Tensor, input_bit_width: Tensor,
+                input_is_signed: bool) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:
+        bit_width = self.msb_clamp_bit_width_impl()
+        int_threshold = self.int_scaling_impl(bit_width)
+        pre_threshold = self.pre_scaling_impl(x, input_bit_width, input_is_signed)
+        pre_scale = pre_threshold / int_threshold
+        pre_zero_point = self.pre_zero_point_impl(x, pre_scale, bit_width)
+        threshold = self.scaling_impl(x)
+        scale = threshold / int_threshold
+        zero_point = self.zero_point_impl(x, scale, bit_width)
+        y = self.decoupled_int_quant(pre_scale, pre_zero_point, scale, zero_point, bit_width, x)
+        return y, scale, zero_point, bit_width, pre_scale, pre_zero_point

 

 
             

diff --git a/docs/_modules/brevitas/core/quant/int_base.html b/docs/_modules/brevitas/core/quant/int_base.html
index 2d620a38a..b14c1feee 100644
--- a/docs/_modules/brevitas/core/quant/int_base.html
+++ b/docs/_modules/brevitas/core/quant/int_base.html
@@ -5,7 +5,7 @@
   
     
     
-    brevitas.core.quant.int_base — Brevitas 0.7.2.dev139+g0c2e90d documentation
+    brevitas.core.quant.int_base — Brevitas 0.9.0 documentation
   
   
   
@@ -350,54 +350,19 @@
             

               
   
Source code for brevitas.core.quant.int_base
-# Copyright (c) 2018-     Xilinx, Inc              (Alessandro Pappalardo)
-# Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)
-# Copyright (c) 2014-     Facebook, Inc            (Soumith Chintala)
-# Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
-# Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)
-# Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
-# Copyright (c) 2011-2013 NYU                      (Clement Farabet)
-# Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)
-# Copyright (c) 2006      Idiap Research Institute (Samy Bengio)
-# Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
-
-# All rights reserved.
-
-# Redistribution and use in source and binary forms, with or without
-# modification, are permitted provided that the following conditions are met:
-
-# 1. Redistributions of source code must retain the above copyright
-#    notice, this list of conditions and the following disclaimer.
-
-# 2. Redistributions in binary form must reproduce the above copyright
-#    notice, this list of conditions and the following disclaimer in the
-#    documentation and/or other materials provided with the distribution.
-
-# 3. Neither the names of Xilinx, Facebook, Deepmind Technologies, NYU,
-#    NEC Laboratories America and IDIAP Research Institute nor the names
-#    of its contributors may be used to endorse or promote products derived
-#    from this software without specific prior written permission.
-
-# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
-# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
-# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
-# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
-# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
-# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
-# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
-# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
-# POSSIBILITY OF SUCH DAMAGE.
+# Copyright (C) 2023, Advanced Micro Devices, Inc. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
 
 import torch
 from torch import Tensor
 from torch.nn import Module
 
 import brevitas
-from brevitas.function.ops import max_int, min_int
-from brevitas.core.function_wrapper import RoundSte, TensorClamp
+from brevitas.core.function_wrapper import RoundSte
+from brevitas.core.function_wrapper import TensorClamp
 from brevitas.core.quant.delay import DelayWrapper
+from brevitas.function.ops import max_int
+from brevitas.function.ops import min_int
 
 
 
[docs]class IntQuant(brevitas.jit.ScriptModule):
@@ -449,12 +414,7 @@ 
Source code for brevitas.core.quant.int_base
self.delay_wrapper = DelayWrapper(quant_delay_steps)
 
 
[docs]    @brevitas.jit.script_method_110_disabled
-    def to_int(
-            self,
-            scale: Tensor,
-            zero_point: Tensor,
-            bit_width: Tensor,
-            x: Tensor) -> Tensor:
+    def to_int(self, scale: Tensor, zero_point: Tensor, bit_width: Tensor, x: Tensor) -> Tensor:
         y = x / scale
         y = y + zero_point
         min_int_val = self.min_int(bit_width)
@@ -472,12 +432,7 @@ 
Source code for brevitas.core.quant.int_base
return max_int(self.signed, self.narrow_range, bit_width)

 
 
[docs]    @brevitas.jit.script_method
-    def forward(
-            self,
-            scale: Tensor,
-            zero_point: Tensor,
-            bit_width: Tensor,
-            x: Tensor) -> Tensor:
+    def forward(self, scale: Tensor, zero_point: Tensor, bit_width: Tensor, x: Tensor) -> Tensor:
         y_int = self.to_int(scale, zero_point, bit_width, x)
         y = y_int - zero_point
         y = y * scale
@@ -533,10 +488,7 @@ 
Source code for brevitas.core.quant.int_base
[docs]    @brevitas.jit.script_method_110_disabled
     def to_int(
-            self,
-            pre_scale: Tensor,
-            pre_zero_point: Tensor,
-            bit_width: Tensor,
+            self, pre_scale: Tensor, pre_zero_point: Tensor, bit_width: Tensor,
             x: Tensor) -> Tensor:
         y = x / pre_scale
         y = y + pre_zero_point
diff --git a/docs/_modules/brevitas/core/quant/ternary.html b/docs/_modules/brevitas/core/quant/ternary.html
index c342cea24..b84e9aaa8 100644
--- a/docs/_modules/brevitas/core/quant/ternary.html
+++ b/docs/_modules/brevitas/core/quant/ternary.html
@@ -5,7 +5,7 @@
   
     
     
-    brevitas.core.quant.ternary — Brevitas 0.7.2.dev139+g0c2e90d documentation
+    brevitas.core.quant.ternary — Brevitas 0.9.0 documentation
   
   
   
@@ -350,45 +350,8 @@
             

               
   
Source code for brevitas.core.quant.ternary
-# Copyright (c) 2018-     Xilinx, Inc              (Alessandro Pappalardo)
-# Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)
-# Copyright (c) 2014-     Facebook, Inc            (Soumith Chintala)
-# Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
-# Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)
-# Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
-# Copyright (c) 2011-2013 NYU                      (Clement Farabet)
-# Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)
-# Copyright (c) 2006      Idiap Research Institute (Samy Bengio)
-# Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
-
-# All rights reserved.
-
-# Redistribution and use in source and binary forms, with or without
-# modification, are permitted provided that the following conditions are met:
-
-# 1. Redistributions of source code must retain the above copyright
-#    notice, this list of conditions and the following disclaimer.
-
-# 2. Redistributions in binary form must reproduce the above copyright
-#    notice, this list of conditions and the following disclaimer in the
-#    documentation and/or other materials provided with the distribution.
-
-# 3. Neither the names of Xilinx, Facebook, Deepmind Technologies, NYU,
-#    NEC Laboratories America and IDIAP Research Institute nor the names
-#    of its contributors may be used to endorse or promote products derived
-#    from this software without specific prior written permission.
-
-# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
-# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
-# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
-# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
-# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
-# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
-# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
-# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
-# POSSIBILITY OF SUCH DAMAGE.
+# Copyright (C) 2023, Advanced Micro Devices, Inc. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
 
 from typing import Tuple
 
@@ -397,10 +360,10 @@ 
Source code for brevitas.core.quant.ternary
 from torch.nn import Module
 
 import brevitas
-from brevitas.function.ops_ste import ternary_sign_ste
 from brevitas.core.bit_width import BitWidthConst
-from brevitas.core.utils import StatelessBuffer
 from brevitas.core.quant.delay import DelayWrapper
+from brevitas.core.utils import StatelessBuffer
+from brevitas.function.ops_ste import ternary_sign_ste
 
 
 
[docs]class TernaryQuant(brevitas.jit.ScriptModule):
diff --git a/docs/_modules/brevitas/core/restrict_val.html b/docs/_modules/brevitas/core/restrict_val.html
index fee8c3c7c..61c262e4e 100644
--- a/docs/_modules/brevitas/core/restrict_val.html
+++ b/docs/_modules/brevitas/core/restrict_val.html
@@ -5,7 +5,7 @@
   
     
     
-    brevitas.core.restrict_val — Brevitas 0.7.2.dev139+g0c2e90d documentation
+    brevitas.core.restrict_val — Brevitas 0.9.0 documentation
   
   
   
@@ -350,58 +350,25 @@
             

               
   
Source code for brevitas.core.restrict_val
-# Copyright (c) 2018-     Xilinx, Inc              (Alessandro Pappalardo)
-# Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)
-# Copyright (c) 2014-     Facebook, Inc            (Soumith Chintala)
-# Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
-# Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)
-# Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
-# Copyright (c) 2011-2013 NYU                      (Clement Farabet)
-# Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)
-# Copyright (c) 2006      Idiap Research Institute (Samy Bengio)
-# Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
-
-# All rights reserved.
-
-# Redistribution and use in source and binary forms, with or without
-# modification, are permitted provided that the following conditions are met:
-
-# 1. Redistributions of source code must retain the above copyright
-#    notice, this list of conditions and the following disclaimer.
-
-# 2. Redistributions in binary form must reproduce the above copyright
-#    notice, this list of conditions and the following disclaimer in the
-#    documentation and/or other materials provided with the distribution.
-
-# 3. Neither the names of Xilinx, Facebook, Deepmind Technologies, NYU,
-#    NEC Laboratories America and IDIAP Research Institute nor the names
-#    of its contributors may be used to endorse or promote products derived
-#    from this software without specific prior written permission.
-
-# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
-# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
-# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
-# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
-# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
-# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
-# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
-# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
-# POSSIBILITY OF SUCH DAMAGE.
-
-from typing import Callable, Union, Optional
+# Copyright (C) 2023, Advanced Micro Devices, Inc. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
 import math
+from typing import Callable, Optional, Union
 
 import torch
 from torch import Tensor
 from torch.nn import Module
 
 import brevitas
-from brevitas.inject.enum import RestrictValueType, FloatToIntImplType  # retrocompatibility
-
-from brevitas.core.function_wrapper import Identity, PowerOfTwo, LogTwo, InplaceLogTwo
-from brevitas.core.function_wrapper import ScalarClampMinSte, RoundSte
+from brevitas.core.function_wrapper import Identity
+from brevitas.core.function_wrapper import InplaceLogTwo
+from brevitas.core.function_wrapper import LogTwo
+from brevitas.core.function_wrapper import PowerOfTwo
+from brevitas.core.function_wrapper import RoundSte
+from brevitas.core.function_wrapper import ScalarClampMinSte
+from brevitas.inject.enum import FloatToIntImplType  # retrocompatibility
+from brevitas.inject.enum import RestrictValueType
 
 assert RestrictValueType  # prevent removal of unused import
 assert FloatToIntImplType
@@ -409,10 +376,7 @@ 
Source code for brevitas.core.restrict_val
 
 class _RestrictClampValue(brevitas.jit.ScriptModule):
 
-    def __init__(
-            self,
-            scaling_min_val: Optional[float],
-            restrict_value_impl: Optional[Module]):
+    def __init__(self, scaling_min_val: Optional[float], restrict_value_impl: Optional[Module]):
         super(_RestrictClampValue, self).__init__()
         if scaling_min_val is not None and scaling_min_val != 0:
             self.clamp_min_ste = ScalarClampMinSte(scaling_min_val)
@@ -432,9 +396,7 @@ 
Source code for brevitas.core.restrict_val
 
 class _RestrictValue(brevitas.jit.ScriptModule):
 
-    def __init__(
-            self,
-            restrict_value_impl: Optional[Module]):
+    def __init__(self, restrict_value_impl: Optional[Module]):
         super(_RestrictValue, self).__init__()
         if restrict_value_impl is not None:
             self.restrict_value_impl = restrict_value_impl
@@ -449,9 +411,7 @@ 
Source code for brevitas.core.restrict_val
 
 class _ClampValue(brevitas.jit.ScriptModule):
 
-    def __init__(
-            self,
-            scaling_min_val: Optional[float]):
+    def __init__(self, scaling_min_val: Optional[float]):
         super(_ClampValue, self).__init__()
         if scaling_min_val is not None and scaling_min_val != 0:
             self.clamp_min_ste = ScalarClampMinSte(scaling_min_val)
diff --git a/docs/_modules/brevitas/core/scaling/int_scaling.html b/docs/_modules/brevitas/core/scaling/int_scaling.html
index bece29ad6..88094781f 100644
--- a/docs/_modules/brevitas/core/scaling/int_scaling.html
+++ b/docs/_modules/brevitas/core/scaling/int_scaling.html
@@ -5,7 +5,7 @@
   
     
     
-    brevitas.core.scaling.int_scaling — Brevitas 0.7.2.dev139+g0c2e90d documentation
+    brevitas.core.scaling.int_scaling — Brevitas 0.9.0 documentation
   
   
   
@@ -350,51 +350,14 @@
             

               
   
Source code for brevitas.core.scaling.int_scaling
-# Copyright (c) 2018-     Xilinx, Inc              (Alessandro Pappalardo)
-# Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)
-# Copyright (c) 2014-     Facebook, Inc            (Soumith Chintala)
-# Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
-# Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)
-# Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
-# Copyright (c) 2011-2013 NYU                      (Clement Farabet)
-# Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)
-# Copyright (c) 2006      Idiap Research Institute (Samy Bengio)
-# Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
-
-# All rights reserved.
-
-# Redistribution and use in source and binary forms, with or without
-# modification, are permitted provided that the following conditions are met:
-
-# 1. Redistributions of source code must retain the above copyright
-#    notice, this list of conditions and the following disclaimer.
-
-# 2. Redistributions in binary form must reproduce the above copyright
-#    notice, this list of conditions and the following disclaimer in the
-#    documentation and/or other materials provided with the distribution.
-
-# 3. Neither the names of Xilinx, Facebook, Deepmind Technologies, NYU,
-#    NEC Laboratories America and IDIAP Research Institute nor the names
-#    of its contributors may be used to endorse or promote products derived
-#    from this software without specific prior written permission.
-
-# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
-# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
-# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
-# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
-# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
-# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
-# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
-# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
-# POSSIBILITY OF SUCH DAMAGE.
-
+# Copyright (C) 2023, Advanced Micro Devices, Inc. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
 
 from torch import Tensor
 
 import brevitas
-from brevitas.function.ops import min_int, max_int
+from brevitas.function.ops import max_int
+from brevitas.function.ops import min_int
 
 
 
[docs]class IntScaling(brevitas.jit.ScriptModule):
@@ -408,7 +371,7 @@ 
Source code for brevitas.core.scaling.int_scaling
[docs]    @brevitas.jit.script_method
     def forward(self, bit_width: Tensor) -> Tensor:
         if self.signed:
-            return - min_int(self.signed, self.narrow_range, bit_width)
+            return -min_int(self.signed, self.narrow_range, bit_width)
         else:
             return max_int(self.signed, self.narrow_range, bit_width)

 
diff --git a/docs/_modules/brevitas/core/scaling/runtime.html b/docs/_modules/brevitas/core/scaling/runtime.html
index 7786e03c3..d6380e09a 100644
--- a/docs/_modules/brevitas/core/scaling/runtime.html
+++ b/docs/_modules/brevitas/core/scaling/runtime.html
@@ -5,7 +5,7 @@
   
     
     
-    brevitas.core.scaling.runtime — Brevitas 0.7.2.dev139+g0c2e90d documentation
+    brevitas.core.scaling.runtime — Brevitas 0.9.0 documentation
   
   
   
@@ -350,58 +350,25 @@
             

               
   
Source code for brevitas.core.scaling.runtime
-# Copyright (c) 2018-     Xilinx, Inc              (Alessandro Pappalardo)
-# Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)
-# Copyright (c) 2014-     Facebook, Inc            (Soumith Chintala)
-# Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
-# Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)
-# Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
-# Copyright (c) 2011-2013 NYU                      (Clement Farabet)
-# Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)
-# Copyright (c) 2006      Idiap Research Institute (Samy Bengio)
-# Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
-
-# All rights reserved.
-
-# Redistribution and use in source and binary forms, with or without
-# modification, are permitted provided that the following conditions are met:
-
-# 1. Redistributions of source code must retain the above copyright
-#    notice, this list of conditions and the following disclaimer.
-
-# 2. Redistributions in binary form must reproduce the above copyright
-#    notice, this list of conditions and the following disclaimer in the
-#    documentation and/or other materials provided with the distribution.
-
-# 3. Neither the names of Xilinx, Facebook, Deepmind Technologies, NYU,
-#    NEC Laboratories America and IDIAP Research Institute nor the names
-#    of its contributors may be used to endorse or promote products derived
-#    from this software without specific prior written permission.
-
-# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
-# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
-# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
-# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
-# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
-# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
-# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
-# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
-# POSSIBILITY OF SUCH DAMAGE.
-
-from typing import Tuple, Optional, List, Union
+# Copyright (C) 2023, Advanced Micro Devices, Inc. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+from typing import List, Optional, Tuple, Union
 
 import torch
-from torch.nn import Parameter, Module
+from torch.nn import Module
+from torch.nn import Parameter
 
 import brevitas
 import brevitas.config as config
 from brevitas.core.function_wrapper import Identity
-from brevitas.function.ops_ste import abs_binary_sign_grad
-
-from brevitas.core.stats import _ParameterListStats, _RuntimeStats, DEFAULT_MOMENTUM
 from brevitas.core.restrict_val import _RestrictClampValue
+from brevitas.core.stats import _ParameterListStats
+from brevitas.core.stats import _RuntimeStats
+from brevitas.core.stats import DEFAULT_MOMENTUM
+from brevitas.core.utils import ParameterWrapper
+from brevitas.core.utils import StatelessBuffer
+from brevitas.function.ops_ste import abs_binary_sign_grad
 
 
 
[docs]class StatsFromParameterScaling(brevitas.jit.ScriptModule):
@@ -415,6 +382,7 @@ 
Source code for brevitas.core.scaling.runtime
restrict_scaling_impl: Module,
             scaling_shape: Tuple[int, ...],
             affine_rescaling: bool = False,
+            affine_shift_scale: bool = False,
             scaling_min_val: Optional[float] = None) -> None:
         super(StatsFromParameterScaling, self).__init__()
         self.parameter_list_stats = _ParameterListStats(
@@ -427,7 +395,8 @@ 
Source code for brevitas.core.scaling.runtime
restrict_scaling_impl,
             scaling_shape,
             scaling_min_val,
-            affine_rescaling)
+            affine_rescaling,
+            affine_shift_scale)
 
 
[docs]    @brevitas.jit.script_method
     def forward(self, ignored: torch.Tensor) -> torch.Tensor:
@@ -441,12 +410,15 @@ 
Source code for brevitas.core.scaling.runtime
self,
             restrict_scaling_impl: Module,
             scaling_shape: Tuple[int, ...],
-            scaling_min_val: Optional[float] = None,
-            affine_rescaling: bool = False) -> None:
+            scaling_min_val: Optional[float],
+            affine_rescaling: bool,
+            affine_shift_scale: bool) -> None:
         super(_StatsScaling, self).__init__()
-
+        if affine_shift_scale and not affine_rescaling:
+            raise RuntimeError(
+                "Disabling shifting of the scale requires to enable affine rescaling first.")
         if affine_rescaling:
-            self.affine_rescaling = _AffineRescaling(scaling_shape)
+            self.affine_rescaling = _AffineRescaling(scaling_shape, affine_shift_scale)
         else:
             self.affine_rescaling = Identity()
         self.restrict_clamp_scaling = _RestrictClampValue(scaling_min_val, restrict_scaling_impl)
@@ -468,7 +440,8 @@ 
Source code for brevitas.core.scaling.runtime
scaling_stats_input_view_shape_impl: Module,
             restrict_scaling_impl: Module,
             scaling_shape: Tuple[int, ...],
-            affine_rescaling: bool,
+            affine_rescaling: bool = False,
+            affine_shift_scale: bool = False,
             scaling_stats_momentum: float = DEFAULT_MOMENTUM,
             scaling_min_val: Optional[float] = None) -> None:
         super(RuntimeStatsScaling, self).__init__()
@@ -482,7 +455,8 @@ 
Source code for brevitas.core.scaling.runtime
restrict_scaling_impl,
             scaling_shape,
             scaling_min_val,
-            affine_rescaling)
+            affine_rescaling,
+            affine_shift_scale)
 
 
[docs]    @brevitas.jit.script_method
     def forward(self, x: torch.Tensor):
@@ -492,19 +466,23 @@ 
Source code for brevitas.core.scaling.runtime
class _AffineRescaling(brevitas.jit.ScriptModule):
 
-    def __init__(self, scaling_shape):
+    def __init__(self, scaling_shape, shift_scale):
         super(_AffineRescaling, self).__init__()
         self.affine_weight = Parameter(torch.ones(scaling_shape))
-        self.affine_bias = Parameter(torch.zeros(scaling_shape))
+        if shift_scale:
+            self.affine_bias = ParameterWrapper(torch.zeros(scaling_shape))
+        else:
+            self.affine_bias = StatelessBuffer(torch.tensor(0.))
 
     @brevitas.jit.script_method
     def forward(self, x):
-        out = x * self.affine_weight + self.affine_bias
+        out = x * self.affine_weight + self.affine_bias()
         out = abs_binary_sign_grad(out)
         return out
 
-    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict,
-                              missing_keys, unexpected_keys, error_msgs):
+    def _load_from_state_dict(
+            self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys,
+            error_msgs):
         super(_AffineRescaling, self)._load_from_state_dict(
             state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
         affine_weight_key = prefix + 'affine_weight'
diff --git a/docs/_modules/brevitas/core/scaling/standalone.html b/docs/_modules/brevitas/core/scaling/standalone.html
index 18c9a98c7..80b2bb9eb 100644
--- a/docs/_modules/brevitas/core/scaling/standalone.html
+++ b/docs/_modules/brevitas/core/scaling/standalone.html
@@ -5,7 +5,7 @@
   
     
     
-    brevitas.core.scaling.standalone — Brevitas 0.7.2.dev139+g0c2e90d documentation
+    brevitas.core.scaling.standalone — Brevitas 0.9.0 documentation
   
   
   
@@ -350,61 +350,31 @@
             

               
   
Source code for brevitas.core.scaling.standalone
-# Copyright (c) 2018-     Xilinx, Inc              (Alessandro Pappalardo)
-# Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)
-# Copyright (c) 2014-     Facebook, Inc            (Soumith Chintala)
-# Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
-# Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)
-# Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
-# Copyright (c) 2011-2013 NYU                      (Clement Farabet)
-# Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)
-# Copyright (c) 2006      Idiap Research Institute (Samy Bengio)
-# Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
-
-# All rights reserved.
-
-# Redistribution and use in source and binary forms, with or without
-# modification, are permitted provided that the following conditions are met:
-
-# 1. Redistributions of source code must retain the above copyright
-#    notice, this list of conditions and the following disclaimer.
-
-# 2. Redistributions in binary form must reproduce the above copyright
-#    notice, this list of conditions and the following disclaimer in the
-#    documentation and/or other materials provided with the distribution.
-
-# 3. Neither the names of Xilinx, Facebook, Deepmind Technologies, NYU,
-#    NEC Laboratories America and IDIAP Research Institute nor the names
-#    of its contributors may be used to endorse or promote products derived
-#    from this software without specific prior written permission.
-
-# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
-# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
-# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
-# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
-# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
-# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
-# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
-# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
-# POSSIBILITY OF SUCH DAMAGE.
+# Copyright (C) 2023, Advanced Micro Devices, Inc. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
 
+from typing import Optional, Tuple, Union
 import warnings
-from typing import Tuple, Optional, Union
 
 import torch
 from torch import Tensor
-from torch.nn import Parameter, Module
+from torch.nn import Module
+from torch.nn import Parameter
 
 import brevitas
 import brevitas.config as config
-from brevitas.function import abs_binary_sign_grad
 from brevitas.core.function_wrapper import Identity
 from brevitas.core.function_wrapper import OverBatchOverTensorView
-from brevitas.core.utils import StatelessBuffer, inplace_momentum_update, inplace_tensor_mul
-from brevitas.core.restrict_val import _RestrictClampValue, _RestrictValue, _ClampValue
-from brevitas.core.stats import _Stats, SCALAR_SHAPE, DEFAULT_MOMENTUM
+from brevitas.core.restrict_val import _ClampValue
+from brevitas.core.restrict_val import _RestrictClampValue
+from brevitas.core.restrict_val import _RestrictValue
+from brevitas.core.stats import _Stats
+from brevitas.core.stats import DEFAULT_MOMENTUM
+from brevitas.core.stats import SCALAR_SHAPE
+from brevitas.core.utils import inplace_momentum_update
+from brevitas.core.utils import inplace_tensor_mul
+from brevitas.core.utils import StatelessBuffer
+from brevitas.function import abs_binary_sign_grad
 
 
 
[docs]class ConstScaling(brevitas.jit.ScriptModule):
@@ -437,6 +407,7 @@ 
Source code for brevitas.core.scaling.standalone

     Note:
         Maps to scaling_impl_type == ScalingImplType.CONST == 'CONST' == 'const' in higher-level APIs.
     """
+
     def __init__(
             self,
             scaling_init: Union[float, Tensor],
@@ -498,6 +469,7 @@ 
Source code for brevitas.core.scaling.standalone

     Note:
         Maps to scaling_impl_type == ScalingImplType.PARAMETER == 'PARAMETER' == 'parameter' in higher-level APIs.
     """
+
     def __init__(
             self,
             scaling_init: Union[float, Tensor],
@@ -506,10 +478,8 @@ 
Source code for brevitas.core.scaling.standalone

             scaling_min_val: Optional[float] = None) -> None:
         super(ParameterScaling, self).__init__()
 
-        if (isinstance(scaling_init, Tensor)
-                and scaling_shape is not None
-                and scaling_init.shape != SCALAR_SHAPE
-                and scaling_init.shape != scaling_shape):
+        if (isinstance(scaling_init, Tensor) and scaling_shape is not None and
+                scaling_init.shape != SCALAR_SHAPE and scaling_init.shape != scaling_shape):
             raise RuntimeError("scaling_init.shape is non-scalar and != from scaling_shape.")
 
         if isinstance(scaling_init, Tensor):
@@ -528,8 +498,9 @@ 
Source code for brevitas.core.scaling.standalone

         value = abs_binary_sign_grad(self.restrict_clamp_scaling(self.value))
         return value

 
-    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict,
-                              missing_keys, unexpected_keys, error_msgs):
+    def _load_from_state_dict(
+            self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys,
+            error_msgs):
         value_key = prefix + 'value'
         retrocomp_value_key = prefix + 'learned_value'
         if retrocomp_value_key in state_dict:  # retrocompatibility
@@ -545,7 +516,7 @@ 
Source code for brevitas.core.scaling.standalone

     ScriptModule implementation of a learned scale factor initialized from runtime statistics.
     The implementation works in two phases. During the first phase, statistics are collected in
     the same fashion as batchnorm, meaning that while the module is in training mode a set of per-batch
-    statistics are computed and returned, while in background an average of them is retained and returned 
+    statistics are computed and returned, while in background an average of them is retained and returned
     in inference mode. During the second phase, the average accumulated during the first
     phase is used to initialize a learned torch.nn.Parameter, and then the behaviour is the same
     as ParameterScaling.
@@ -585,7 +556,7 @@ 
Source code for brevitas.core.scaling.standalone

         Maps to scaling_impl_type == ScalingImplType.PARAMETER_FROM_STATS == 'PARAMETER_FROM_STATS'
         == 'parameter_from_stats' when applied to runtime values (inputs/outputs/activations) in higher-level APIs.
     """
-    __constants__ = ['collect_stats_steps', 'momentum']
+    __constants__ = ['momentum']
 
     def __init__(
             self,
@@ -598,7 +569,7 @@ 
Source code for brevitas.core.scaling.standalone

             scaling_min_val: Optional[float] = None) -> None:
         super(ParameterFromRuntimeStatsScaling, self).__init__()
         assert collect_stats_steps > 0, 'Steps should be more than 0'
-        self.collect_stats_steps = collect_stats_steps
+        self.collect_stats_steps: int = brevitas.jit.Attribute(collect_stats_steps, int)
         self.counter: int = brevitas.jit.Attribute(0, int)
         self.stats_input_view_shape_impl = scaling_stats_input_view_shape_impl
         self.stats = _Stats(scaling_stats_impl, scaling_shape)
@@ -613,14 +584,14 @@ 
Source code for brevitas.core.scaling.standalone

         else:
             self.restrict_inplace_preprocess = Identity()
             self.restrict_preprocess = Identity()
-    
+
 
[docs]    @brevitas.jit.script_method
     def training_forward(self, stats_input: Tensor) -> Tensor:
         if self.counter < self.collect_stats_steps:
             stats_input = self.stats_input_view_shape_impl(stats_input)
             stats = self.stats(stats_input)
             # workaround to avoid find_ununsed_parameter=True in DDP
-            stats = stats + 0. * self.value # stats gradient will change from None to 0.
+            stats = stats + 0. * self.value  # stats gradient will change from None to 0.
             clamped_stats = self.clamp_scaling(stats)
             new_counter = self.counter + 1
             if self.counter == 0:
@@ -653,7 +624,7 @@ 
Source code for brevitas.core.scaling.standalone

 
 
[docs]    def state_dict(self, destination=None, prefix='', keep_vars=False):
         output_dict = super(ParameterFromRuntimeStatsScaling, self).state_dict(
-            destination, prefix, keep_vars)        
+            destination=destination, prefix=prefix, keep_vars=keep_vars)
         # Avoid saving the buffer
         del output_dict[prefix + 'buffer']
         # Avoid saving the init value
@@ -664,17 +635,19 @@ 
Source code for brevitas.core.scaling.standalone

             output_dict[prefix + 'value'] = self.restrict_preprocess(self.buffer)
         return output_dict

 
-    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict,
-                              missing_keys, unexpected_keys, error_msgs):
-        super(ParameterFromRuntimeStatsScaling, self)._load_from_state_dict(
-            state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
-        value_key = prefix + 'value'
-        # Buffer is supposed to be always missing
-        missing_keys.remove(prefix + 'buffer')
+    def _load_from_state_dict(
+            self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys,
+            error_msgs):
         # Retrocompatibility with older ParameterScaling, for when scaling impl is switched over
+        value_key = prefix + 'value'
         retrocomp_value_key = prefix + 'learned_value'
         if retrocomp_value_key in state_dict:
             state_dict[value_key] = state_dict.pop(retrocomp_value_key)
+
+        super(ParameterFromRuntimeStatsScaling, self)._load_from_state_dict(
+            state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
+        # Buffer is supposed to be always missing
+        missing_keys.remove(prefix + 'buffer')
         # Pytorch stores training flag as a buffer with JIT enabled
         training_key = prefix + 'training'
         if training_key in missing_keys:
diff --git a/docs/_modules/brevitas/core/stats/stats_op.html b/docs/_modules/brevitas/core/stats/stats_op.html
index 85fb0859e..079463f09 100644
--- a/docs/_modules/brevitas/core/stats/stats_op.html
+++ b/docs/_modules/brevitas/core/stats/stats_op.html
@@ -5,7 +5,7 @@
   
     
     
-    brevitas.core.stats.stats_op — Brevitas 0.7.2.dev139+g0c2e90d documentation
+    brevitas.core.stats.stats_op — Brevitas 0.9.0 documentation
   
   
   
@@ -350,60 +350,22 @@
             

               
   
Source code for brevitas.core.stats.stats_op
-# Copyright (c) 2018-     Xilinx, Inc              (Alessandro Pappalardo)
-# Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)
-# Copyright (c) 2014-     Facebook, Inc            (Soumith Chintala)
-# Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
-# Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)
-# Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
-# Copyright (c) 2011-2013 NYU                      (Clement Farabet)
-# Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)
-# Copyright (c) 2006      Idiap Research Institute (Samy Bengio)
-# Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
-
-# All rights reserved.
-
-# Redistribution and use in source and binary forms, with or without
-# modification, are permitted provided that the following conditions are met:
-
-# 1. Redistributions of source code must retain the above copyright
-#    notice, this list of conditions and the following disclaimer.
-
-# 2. Redistributions in binary form must reproduce the above copyright
-#    notice, this list of conditions and the following disclaimer in the
-#    documentation and/or other materials provided with the distribution.
-
-# 3. Neither the names of Xilinx, Facebook, Deepmind Technologies, NYU,
-#    NEC Laboratories America and IDIAP Research Institute nor the names
-#    of its contributors may be used to endorse or promote products derived
-#    from this software without specific prior written permission.
-
-# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
-# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
-# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
-# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
-# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
-# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
-# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
-# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
-# POSSIBILITY OF SUCH DAMAGE.
+# Copyright (C) 2023, Advanced Micro Devices, Inc. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
 
-
-from typing import Optional, Tuple
 import math
+from typing import Optional, Tuple
 
 import torch
 from torch import Tensor
 from torch.nn import Parameter
 
 import brevitas
+from brevitas import config
 from brevitas.core.utils import StatelessBuffer
 from brevitas.function.ops import max_int
-from brevitas import config
-from .stats_wrapper import SCALAR_SHAPE
 
+from .stats_wrapper import SCALAR_SHAPE
 
 DEFAULT_STD_DEV_EPSILON = 1e-8
 
@@ -430,7 +392,8 @@ 
Source code for brevitas.core.stats.stats_op
[docs]class AbsPercentile(brevitas.jit.ScriptModule):
     __constants__ = ['q', 'stats_reduce_dim']
 
-    def __init__(self, high_percentile_q: float, stats_reduce_dim: Optional[int], percentile_q = None):
+    def __init__(
+            self, high_percentile_q: float, stats_reduce_dim: Optional[int], percentile_q=None):
         super(AbsPercentile, self).__init__()
         if percentile_q is not None:
             raise RuntimeError("percentile_q is deprecated, please pass high_percentile_q.")
@@ -453,8 +416,8 @@ 
Source code for brevitas.core.stats.stats_op
k = int(math.floor(.01 * self.q * dim_slice.numel() + 0.5))
             result = x.abs().kthvalue(k, dim=self.stats_reduce_dim).values
         return result

-    
-    
+
+
 
[docs]class NegativePercentileOrZero(brevitas.jit.ScriptModule):
     __constants__ = ['stats_reduce_dim', 'q']
 
@@ -481,12 +444,16 @@ 
Source code for brevitas.core.stats.stats_op
result = torch.where(
             result <= self.zero().to(result.dtype), result, self.zero().to(result.dtype))
         return result

-    
-    
+
+
 
[docs]class PercentileInterval(brevitas.jit.ScriptModule):
     __constants__ = ['stats_reduce_dim', 'low_q', 'high_q']
 
-    def __init__(self, low_percentile_q, high_percentile_q, stats_reduce_dim: Optional[int] = None) -> None:
+    def __init__(
+            self,
+            low_percentile_q,
+            high_percentile_q,
+            stats_reduce_dim: Optional[int] = None) -> None:
         super(PercentileInterval, self).__init__()
         self.stats_reduce_dim = stats_reduce_dim
         self.low_q = low_percentile_q
@@ -653,8 +620,9 @@ 
Source code for brevitas.core.stats.stats_op
out = self.impl(x, sigma)
         return out

 
-    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict,
-                              missing_keys, unexpected_keys, error_msgs):
+    def _load_from_state_dict(
+            self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys,
+            error_msgs):
         value_key = prefix + 'sigma'
         retrocomp_value_key = prefix + 'learned_sigma'
         if retrocomp_value_key in state_dict:  # retrocompatibility
@@ -672,7 +640,7 @@ 
Source code for brevitas.core.stats.stats_op
    https://github.com/apache/incubator-mxnet/blob/master/python/mxnet/contrib/quantization.py
     """
 
-    def __init__(self, signed, bit_width_impl, num_bins = 1000 + 1, smoothing_eps=0.0001):
+    def __init__(self, signed, bit_width_impl, num_bins=1000 + 1, smoothing_eps=0.0001):
         super(KLMinimizerThreshold, self).__init__()
         self.num_bins = num_bins
         self.smoothing_eps = smoothing_eps
@@ -738,6 +706,42 @@ 
Source code for brevitas.core.stats.stats_op
min_divergence_idx = torch.argmin(divergence)
         opt_threshold = thresholds[min_divergence_idx]
         return opt_threshold

+
+
+
[docs]class L1Norm(brevitas.jit.ScriptModule):
+    """ScriptModule implementation to collect per-channel L1 normalization stats
+    for weight normalization-based quantization."""
+    __constants__ = ['stats_reduce_dim']
+
+    def __init__(self, stats_reduce_dim: Optional[int] = None) -> None:
+        super(L1Norm, self).__init__()
+        self.stats_reduce_dim = stats_reduce_dim
+
+
[docs]    @brevitas.jit.script_method
+    def forward(self, x: Tensor):
+        if self.stats_reduce_dim is None:
+            # Need to be able to return the max per-channel L1 norm as a scalar
+            raise NotImplementedError("L1 normalization is not supported per-tensor yet.")
+        else:
+            return x.norm(p=1, dim=self.stats_reduce_dim, keepdim=True)

+
+
+
[docs]class L2Norm(brevitas.jit.ScriptModule):
+    """ScriptModule implementation to collect per-channel L2 normalization stats
+    for weight normalization-based quantization."""
+    __constants__ = ['stats_reduce_dim']
+
+    def __init__(self, stats_reduce_dim: Optional[int] = None) -> None:
+        super(L2Norm, self).__init__()
+        self.stats_reduce_dim = stats_reduce_dim
+
+
[docs]    @brevitas.jit.script_method
+    def forward(self, x: Tensor):
+        if self.stats_reduce_dim is None:
+            # Need to be able to return the max per-channel L2 norm as a scalar
+            raise NotImplementedError("L2 normalization is not supported per-tensor yet.")
+        else:
+            return x.norm(p=2, dim=self.stats_reduce_dim, keepdim=True)

 

 
             

diff --git a/docs/_modules/brevitas/core/utils.html b/docs/_modules/brevitas/core/utils.html
index 14e8c63a4..ae6d2162f 100644
--- a/docs/_modules/brevitas/core/utils.html
+++ b/docs/_modules/brevitas/core/utils.html
@@ -5,7 +5,7 @@
   
     
     
-    brevitas.core.utils — Brevitas 0.7.2.dev139+g0c2e90d documentation
+    brevitas.core.utils — Brevitas 0.9.0 documentation
   
   
   
@@ -350,7 +350,11 @@
             

               
   
Source code for brevitas.core.utils
-from typing import Optional
+# Copyright (C) 2023, Advanced Micro Devices, Inc. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+from typing import Optional
+
 import torch
 
 import brevitas
@@ -391,13 +395,14 @@ 
Source code for brevitas.core.utils
     def __init__(self, value: torch.Tensor):
         super(StatelessBuffer, self).__init__()
         self.register_buffer(VALUE_ATTR_NAME, value)
-    
+
 
[docs]    @brevitas.jit.script_method
     def forward(self):
         return self.value.detach()

 
-    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict,
-                              missing_keys, unexpected_keys, error_msgs):
+    def _load_from_state_dict(
+            self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys,
+            error_msgs):
         super(StatelessBuffer, self)._load_from_state_dict(
             state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
         value_key = prefix + VALUE_ATTR_NAME
@@ -405,9 +410,32 @@ 
Source code for brevitas.core.utils
             missing_keys.remove(value_key)
 
 
[docs]    def state_dict(self, destination=None, prefix='', keep_vars=False):
-        output_dict = super(StatelessBuffer, self).state_dict(destination, prefix, keep_vars)
+        output_dict = super(StatelessBuffer, self).state_dict(
+            destination=destination, prefix=prefix, keep_vars=keep_vars)
         del output_dict[prefix + VALUE_ATTR_NAME]
         return output_dict

+
+
+
[docs]class SingleArgStatelessBuffer(brevitas.jit.ScriptModule):
+
+    def __init__(self, value: torch.Tensor):
+        super(SingleArgStatelessBuffer, self).__init__()
+        self.const = StatelessBuffer(torch.tensor(value))
+
+
[docs]    @brevitas.jit.script_method
+    def forward(self, placeholder):
+        return self.const()

+
+
+
[docs]class ParameterWrapper(brevitas.jit.ScriptModule):
+
+    def __init__(self, value: torch.Tensor):
+        super(ParameterWrapper, self).__init__()
+        self.register_parameter(VALUE_ATTR_NAME, value)
+
+
[docs]    @brevitas.jit.script_method
+    def forward(self):
+        return self.value

 

 
             

diff --git a/docs/_modules/brevitas/core/zero_point.html b/docs/_modules/brevitas/core/zero_point.html
index 47bd5668d..d08ec6d67 100644
--- a/docs/_modules/brevitas/core/zero_point.html
+++ b/docs/_modules/brevitas/core/zero_point.html
@@ -5,7 +5,7 @@
   
     
     
-    brevitas.core.zero_point — Brevitas 0.7.2.dev139+g0c2e90d documentation
+    brevitas.core.zero_point — Brevitas 0.9.0 documentation
   
   
   
@@ -350,66 +350,32 @@
             

               
   
Source code for brevitas.core.zero_point
-# Copyright (c) 2018-     Xilinx, Inc              (Alessandro Pappalardo)
-# Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)
-# Copyright (c) 2014-     Facebook, Inc            (Soumith Chintala)
-# Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
-# Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)
-# Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
-# Copyright (c) 2011-2013 NYU                      (Clement Farabet)
-# Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)
-# Copyright (c) 2006      Idiap Research Institute (Samy Bengio)
-# Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
-
-# All rights reserved.
-
-# Redistribution and use in source and binary forms, with or without
-# modification, are permitted provided that the following conditions are met:
-
-# 1. Redistributions of source code must retain the above copyright
-#    notice, this list of conditions and the following disclaimer.
-
-# 2. Redistributions in binary form must reproduce the above copyright
-#    notice, this list of conditions and the following disclaimer in the
-#    documentation and/or other materials provided with the distribution.
-
-# 3. Neither the names of Xilinx, Facebook, Deepmind Technologies, NYU,
-#    NEC Laboratories America and IDIAP Research Institute nor the names
-#    of its contributors may be used to endorse or promote products derived
-#    from this software without specific prior written permission.
-
-# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
-# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
-# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
-# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
-# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
-# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
-# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
-# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
-# POSSIBILITY OF SUCH DAMAGE.
-
-from typing import Optional, Tuple, Union, List
+# Copyright (C) 2023, Advanced Micro Devices, Inc. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+from typing import List, Optional, Tuple, Union
 
 import torch
 from torch import Tensor
-from torch.nn import Module, Parameter
+from torch.nn import Module
+from torch.nn import Parameter
 
 import brevitas
 from brevitas import config
-from brevitas.core.stats import SCALAR_SHAPE, DEFAULT_MOMENTUM, _ParameterListStats
+from brevitas.core.stats import _ParameterListStats
+from brevitas.core.stats import DEFAULT_MOMENTUM
+from brevitas.core.stats import SCALAR_SHAPE
 from brevitas.function import abs_binary_sign_grad
 
-from .utils import StatelessBuffer, inplace_tensor_add, inplace_momentum_update
-
+from .utils import inplace_momentum_update
+from .utils import inplace_tensor_add
+from .utils import StatelessBuffer
 
 __all__ = [
     'ZeroZeroPoint',
-    'MinUintZeroPoint',
-    'ParameterFromRuntimeMinZeroPoint',
-    'ParameterZeroPoint'
-]
+    'StatsFromParameterZeroPoint',
+    'ParameterFromRuntimeZeroPoint',
+    'ParameterZeroPoint']
 
 
 
[docs]class ZeroZeroPoint(brevitas.jit.ScriptModule):
@@ -426,8 +392,7 @@ 
Source code for brevitas.core.zero_point
 class _ScaleShiftZeroPoint(brevitas.jit.ScriptModule):
     __constants__ = ['quantize_zero_point']
 
-    def __init__(
-            self, int_quant: Module, quantize_zero_point: bool) -> None:
+    def __init__(self, int_quant: Module, quantize_zero_point: bool) -> None:
         super(_ScaleShiftZeroPoint, self).__init__()
         self.int_quant = int_quant
         self.quantize_zero_point = quantize_zero_point
@@ -440,9 +405,9 @@ 
Source code for brevitas.core.zero_point
         else:
             out = zero_point / scale + min_int
         return out
-    
-    
-class StatsFromParameterZeroPoint(brevitas.jit.ScriptModule):
+
+
+
[docs]class StatsFromParameterZeroPoint(brevitas.jit.ScriptModule):
 
     def __init__(
             self,
@@ -460,20 +425,16 @@ 
Source code for brevitas.core.zero_point
             zero_point_stats_input_view_shape_impl,
             zero_point_stats_input_concat_dim,
             tracked_parameter_list)
-        self.scale_shift_zero_point = _ScaleShiftZeroPoint(
-            int_quant, quantize_zero_point)
+        self.scale_shift_zero_point = _ScaleShiftZeroPoint(int_quant, quantize_zero_point)
 
-    @brevitas.jit.script_method
+
[docs]    @brevitas.jit.script_method
     def forward(self, x: Tensor, scale: Tensor, bit_width: Tensor) -> torch.Tensor:
         stats = self.parameter_list_stats()
-        return self.scale_shift_zero_point(- stats, scale, bit_width)
+        return self.scale_shift_zero_point(-stats, scale, bit_width)

 
 
-class ParameterFromRuntimeZeroPoint(brevitas.jit.ScriptModule):
-    __constants__ = ['stats_permute_dims',
-                     'collect_stats_steps',
-                     'zero_point_shape',
-                     'momentum']
+
[docs]class ParameterFromRuntimeZeroPoint(brevitas.jit.ScriptModule):
+    __constants__ = ['stats_permute_dims', 'collect_stats_steps', 'zero_point_shape', 'momentum']
 
     def __init__(
             self,
@@ -494,10 +455,9 @@ 
Source code for brevitas.core.zero_point
         self.value = Parameter(torch.full(zero_point_shape, 0.0))
         self.register_buffer('buffer', torch.full(zero_point_shape, 0.0))
         self.zero_point_stats_impl = zero_point_stats_impl
-        self.scale_shift_zero_point = _ScaleShiftZeroPoint(
-            int_quant, quantize_zero_point)
+        self.scale_shift_zero_point = _ScaleShiftZeroPoint(int_quant, quantize_zero_point)
 
-    @brevitas.jit.script_method
+
[docs]    @brevitas.jit.script_method
     def training_forward(self, x) -> Tensor:
         if self.counter < self.collect_stats_steps:
             stats_input = self.stats_input_view_shape_impl(x)
@@ -518,9 +478,9 @@ 
Source code for brevitas.core.zero_point
             out = self.value
         else:
             out = self.value
-        return out
+        return out

 
-    @brevitas.jit.script_method
+
[docs]    @brevitas.jit.script_method
     def forward(self, x: Tensor, scale: Tensor, bit_width: Tensor) -> Tensor:
         if self.training:
             out = self.training_forward(x)
@@ -531,11 +491,11 @@ 
Source code for brevitas.core.zero_point
                 out = self.value
         out = abs_binary_sign_grad(out)
         out = self.scale_shift_zero_point(out, scale, bit_width)
-        return out
+        return out

 
-    def state_dict(self, destination=None, prefix='', keep_vars=False):
-        output_dict = super(ParameterFromRuntimeZeroPoint, self).state_dict(
-            destination, prefix, keep_vars)        
+
[docs]    def state_dict(self, destination=None, prefix='', keep_vars=False):
+        output_dict = super(ParameterFromRuntimeZeroPoint,
+                            self).state_dict(destination, prefix, keep_vars)
         # Avoid saving the buffer
         del output_dict[prefix + 'buffer']
         # Avoid saving the init value
@@ -544,10 +504,11 @@ 
Source code for brevitas.core.zero_point
         # Save buffer into value for any non-zero number of collection steps
         elif self.counter <= self.collect_stats_steps:
             output_dict[prefix + 'value'] = self.buffer
-        return output_dict
+        return output_dict

 
-    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict,
-                              missing_keys, unexpected_keys, error_msgs):
+    def _load_from_state_dict(
+            self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys,
+            error_msgs):
         super(ParameterFromRuntimeZeroPoint, self)._load_from_state_dict(
             state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
         value_key = prefix + 'value'
@@ -562,13 +523,11 @@ 
Source code for brevitas.core.zero_point
         if value_key not in missing_keys:
             self.counter = self.collect_stats_steps + 1
         if config.IGNORE_MISSING_KEYS and value_key in missing_keys:
-            missing_keys.remove(value_key)
+            missing_keys.remove(value_key)

 
 
 
[docs]class ParameterZeroPoint(brevitas.jit.ScriptModule):
-    __constants__ = ['stats_permute_dims',
-                     'collect_stats_steps',
-                     'momentum']
+    __constants__ = ['stats_permute_dims', 'collect_stats_steps', 'momentum']
 
     def __init__(
             self,
@@ -577,10 +536,9 @@ 
Source code for brevitas.core.zero_point
             quantize_zero_point: bool,
             zero_point_shape: Tuple[int, ...] = None) -> None:
         super(ParameterZeroPoint, self).__init__()
-        if (isinstance(zero_point_init, Tensor)
-                and zero_point_shape is not None
-                and zero_point_init.shape != SCALAR_SHAPE
-                and zero_point_init.shape != zero_point_shape):
+        if (isinstance(zero_point_init, Tensor) and zero_point_shape is not None and
+                zero_point_init.shape != SCALAR_SHAPE and
+                zero_point_init.shape != zero_point_shape):
             raise RuntimeError("zero_point_init.shape is non-scalar and != from zero_point_shape.")
 
         if isinstance(zero_point_init, Tensor):
@@ -590,8 +548,7 @@ 
Source code for brevitas.core.zero_point
         if zero_point_init.shape == SCALAR_SHAPE and zero_point_shape is not None:
             zero_point_init = torch.full(zero_point_shape, zero_point_init)
         self.value = Parameter(zero_point_init)
-        self.scale_shift_zero_point = _ScaleShiftZeroPoint(
-            int_quant, quantize_zero_point)
+        self.scale_shift_zero_point = _ScaleShiftZeroPoint(int_quant, quantize_zero_point)
 
 
[docs]    @brevitas.jit.script_method
     def forward(self, x: Tensor, scale: Tensor, bit_width: Tensor) -> Tensor:
@@ -599,14 +556,14 @@ 
Source code for brevitas.core.zero_point
         out = self.scale_shift_zero_point(out, scale, bit_width)
         return out

 
-    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict,
-                              missing_keys, unexpected_keys, error_msgs):
+    def _load_from_state_dict(
+            self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys,
+            error_msgs):
         super(ParameterZeroPoint, self)._load_from_state_dict(
             state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
         value_key = prefix + 'value'
         if config.IGNORE_MISSING_KEYS and value_key in missing_keys:
             missing_keys.remove(value_key)

-
 

 
             

diff --git a/docs/_modules/brevitas/function/ops.html b/docs/_modules/brevitas/function/ops.html
index e820f0b47..e88ad5104 100644
--- a/docs/_modules/brevitas/function/ops.html
+++ b/docs/_modules/brevitas/function/ops.html
@@ -5,7 +5,7 @@
   
     
     
-    brevitas.function.ops — Brevitas 0.7.2.dev139+g0c2e90d documentation
+    brevitas.function.ops — Brevitas 0.9.0 documentation
   
   
   
@@ -350,45 +350,8 @@
             

               
   
Source code for brevitas.function.ops
-# Copyright (c) 2018-     Xilinx, Inc              (Alessandro Pappalardo)
-# Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)
-# Copyright (c) 2014-     Facebook, Inc            (Soumith Chintala)
-# Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
-# Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)
-# Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
-# Copyright (c) 2011-2013 NYU                      (Clement Farabet)
-# Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)
-# Copyright (c) 2006      Idiap Research Institute (Samy Bengio)
-# Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
-
-# All rights reserved.
-
-# Redistribution and use in source and binary forms, with or without
-# modification, are permitted provided that the following conditions are met:
-
-# 1. Redistributions of source code must retain the above copyright
-#    notice, this list of conditions and the following disclaimer.
-
-# 2. Redistributions in binary form must reproduce the above copyright
-#    notice, this list of conditions and the following disclaimer in the
-#    documentation and/or other materials provided with the distribution.
-
-# 3. Neither the names of Xilinx, Facebook, Deepmind Technologies, NYU,
-#    NEC Laboratories America and IDIAP Research Institute nor the names
-#    of its contributors may be used to endorse or promote products derived
-#    from this software without specific prior written permission.
-
-# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
-# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
-# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
-# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
-# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
-# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
-# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
-# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
-# POSSIBILITY OF SUCH DAMAGE.
+# Copyright (C) 2023, Advanced Micro Devices, Inc. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
 
 """
 Implementation of various core operations often performed as part of quantization.
@@ -571,9 +534,9 @@ 
Source code for brevitas.function.ops
         tensor(0)
     """
     if signed and narrow_range:
-        value = - (2 ** (bit_width - 1)) + 1
+        value = -(2 ** (bit_width - 1)) + 1
     elif signed and not narrow_range:
-        value = - (2 ** (bit_width - 1))
+        value = -(2 ** (bit_width - 1))
     else:
         value = 0 * bit_width
     return value

diff --git a/docs/_modules/brevitas/function/ops_ste.html b/docs/_modules/brevitas/function/ops_ste.html
index febf8fbc7..cca5f5691 100644
--- a/docs/_modules/brevitas/function/ops_ste.html
+++ b/docs/_modules/brevitas/function/ops_ste.html
@@ -5,7 +5,7 @@
   
     
     
-    brevitas.function.ops_ste — Brevitas 0.7.2.dev139+g0c2e90d documentation
+    brevitas.function.ops_ste — Brevitas 0.9.0 documentation
   
   
   
@@ -350,45 +350,8 @@
             

               
   
Source code for brevitas.function.ops_ste
-# Copyright (c) 2019-     Xilinx, Inc              (Giuseppe Franco)
-# Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)
-# Copyright (c) 2014-     Facebook, Inc            (Soumith Chintala)
-# Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
-# Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)
-# Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
-# Copyright (c) 2011-2013 NYU                      (Clement Farabet)
-# Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)
-# Copyright (c) 2006      Idiap Research Institute (Samy Bengio)
-# Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
-
-# All rights reserved.
-
-# Redistribution and use in source and binary forms, with or without
-# modification, are permitted provided that the following conditions are met:
-
-# 1. Redistributions of source code must retain the above copyright
-#    notice, this list of conditions and the following disclaimer.
-
-# 2. Redistributions in binary form must reproduce the above copyright
-#    notice, this list of conditions and the following disclaimer in the
-#    documentation and/or other materials provided with the distribution.
-
-# 3. Neither the names of Xilinx, Facebook, Deepmind Technologies, NYU,
-#    NEC Laboratories America and IDIAP Research Institute nor the names
-#    of its contributors may be used to endorse or promote products derived
-#    from this software without specific prior written permission.
-
-# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
-# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
-# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
-# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
-# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
-# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
-# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
-# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
-# POSSIBILITY OF SUCH DAMAGE.
+# Copyright (C) 2023, Advanced Micro Devices, Inc. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
 
 """
 Implementation of various functions with a straight-through gradient estimators, dispatched to
@@ -404,8 +367,11 @@ 
Source code for brevitas.function.ops_ste
 from torch import Tensor
 
 import brevitas
-from brevitas.function.ops import dpu_round, tensor_clamp
-from brevitas.function.ops import tensor_clamp_, binary_sign, round_to_zero
+from brevitas.function.ops import binary_sign
+from brevitas.function.ops import dpu_round
+from brevitas.function.ops import round_to_zero
+from brevitas.function.ops import tensor_clamp
+from brevitas.function.ops import tensor_clamp_
 
 __all__ = [
     'round_ste',
@@ -419,12 +385,10 @@ 
Source code for brevitas.function.ops_ste
     'ternary_sign_ste',
     'round_to_zero_ste',
     'dpu_round_ste',
-    'abs_binary_sign_grad'
-]
-
+    'abs_binary_sign_grad']
 
 if brevitas.NATIVE_STE_BACKEND_LOADED:
-    fn_prefix = torch 
+    fn_prefix = torch
     script_flag = brevitas.jit.script
 else:
     fn_prefix = brevitas
@@ -509,9 +473,9 @@ 
Source code for brevitas.function.ops_ste
     Function that implements :func:`~brevitas.function.ops.tensor_clamp` with a straight-through
     gradient estimator for the gradient of y w.r.t. to x, while the gradient of y w.r.t. to min_val
     and max_val is always None.
-    
+
     Notes:
-        Wrapper for either :func:`~brevitas.ops.autograd_ste_ops.tensor_clamp_ste_impl` (with 
+        Wrapper for either :func:`~brevitas.ops.autograd_ste_ops.tensor_clamp_ste_impl` (with
         env ``BREVITAS_JIT=0``) or its native just-in-time compiled variant (with
         ``BREVITAS_JIT=1``).
 
@@ -740,7 +704,7 @@ 
Source code for brevitas.function.ops_ste
 
     Notes:
         Wrapper for either :func:`~brevitas.ops.autograd_ste_ops.abs_binary_sign_grad_impl`
-        (with env ``BREVITAS_JIT=0``) or its native just-in-time compiled variant (with 
+        (with env ``BREVITAS_JIT=0``) or its native just-in-time compiled variant (with
         ``BREVITAS_JIT=1``).
 
     Examples:
diff --git a/docs/_modules/brevitas/function/shape.html b/docs/_modules/brevitas/function/shape.html
index 7734520f8..a921b300e 100644
--- a/docs/_modules/brevitas/function/shape.html
+++ b/docs/_modules/brevitas/function/shape.html
@@ -5,7 +5,7 @@
   
     
     
-    brevitas.function.shape — Brevitas 0.7.2.dev139+g0c2e90d documentation
+    brevitas.function.shape — Brevitas 0.9.0 documentation
   
   
   
@@ -350,45 +350,8 @@
             

               
   
Source code for brevitas.function.shape
-# Copyright (c) 2019-     Xilinx, Inc              (Giuseppe Franco)
-# Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)
-# Copyright (c) 2014-     Facebook, Inc            (Soumith Chintala)
-# Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
-# Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)
-# Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
-# Copyright (c) 2011-2013 NYU                      (Clement Farabet)
-# Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)
-# Copyright (c) 2006      Idiap Research Institute (Samy Bengio)
-# Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
-
-# All rights reserved.
-
-# Redistribution and use in source and binary forms, with or without
-# modification, are permitted provided that the following conditions are met:
-
-# 1. Redistributions of source code must retain the above copyright
-#    notice, this list of conditions and the following disclaimer.
-
-# 2. Redistributions in binary form must reproduce the above copyright
-#    notice, this list of conditions and the following disclaimer in the
-#    documentation and/or other materials provided with the distribution.
-
-# 3. Neither the names of Xilinx, Facebook, Deepmind Technologies, NYU,
-#    NEC Laboratories America and IDIAP Research Institute nor the names
-#    of its contributors may be used to endorse or promote products derived
-#    from this software without specific prior written permission.
-
-# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
-# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
-# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
-# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
-# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
-# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
-# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
-# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
-# POSSIBILITY OF SUCH DAMAGE.
+# Copyright (C) 2023, Advanced Micro Devices, Inc. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
 
 """
 Implementation of various functions to compute shapes that induce flattening along certain
@@ -406,8 +369,7 @@ 
Source code for brevitas.function.shape
     'over_tensor',
     'over_output_channels',
     'over_batch_over_tensor',
-    'over_batch_over_output_channels'
-]
+    'over_batch_over_output_channels']
 
 
 
[docs]@brevitas.jit.script
diff --git a/docs/_modules/brevitas/ops/autograd_ste_ops.html b/docs/_modules/brevitas/ops/autograd_ste_ops.html
index 68722280f..92c001b5a 100644
--- a/docs/_modules/brevitas/ops/autograd_ste_ops.html
+++ b/docs/_modules/brevitas/ops/autograd_ste_ops.html
@@ -5,7 +5,7 @@
   
     
     
-    brevitas.ops.autograd_ste_ops — Brevitas 0.7.2.dev139+g0c2e90d documentation
+    brevitas.ops.autograd_ste_ops — Brevitas 0.9.0 documentation
   
   
   
@@ -350,45 +350,8 @@
             

               
   
Source code for brevitas.ops.autograd_ste_ops
-# Copyright (c) 2019-     Xilinx, Inc              (Giuseppe Franco)
-# Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)
-# Copyright (c) 2014-     Facebook, Inc            (Soumith Chintala)
-# Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
-# Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)
-# Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
-# Copyright (c) 2011-2013 NYU                      (Clement Farabet)
-# Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)
-# Copyright (c) 2006      Idiap Research Institute (Samy Bengio)
-# Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
-
-# All rights reserved.
-
-# Redistribution and use in source and binary forms, with or without
-# modification, are permitted provided that the following conditions are met:
-
-# 1. Redistributions of source code must retain the above copyright
-#    notice, this list of conditions and the following disclaimer.
-
-# 2. Redistributions in binary form must reproduce the above copyright
-#    notice, this list of conditions and the following disclaimer in the
-#    documentation and/or other materials provided with the distribution.
-
-# 3. Neither the names of Xilinx, Facebook, Deepmind Technologies, NYU,
-#    NEC Laboratories America and IDIAP Research Institute nor the names
-#    of its contributors may be used to endorse or promote products derived
-#    from this software without specific prior written permission.
-
-# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
-# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
-# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
-# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
-# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
-# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
-# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
-# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
-# POSSIBILITY OF SUCH DAMAGE.
+# Copyright (C) 2023, Advanced Micro Devices, Inc. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
 
 """
 Implementation of various torch.autograd.Function with straight-through estimators.
@@ -397,10 +360,14 @@ 
Source code for brevitas.ops.autograd_ste_ops
from typing import Tuple
 
 import torch
-from torch.autograd import Function
 from torch import Tensor
+from torch.autograd import Function
 
-from brevitas.function.ops import tensor_clamp, binary_sign, round_to_zero, tensor_clamp_, dpu_round
+from brevitas.function.ops import binary_sign
+from brevitas.function.ops import dpu_round
+from brevitas.function.ops import round_to_zero
+from brevitas.function.ops import tensor_clamp
+from brevitas.function.ops import tensor_clamp_
 
 __all__ = [
     'ScalarClampSteFn',
@@ -425,8 +392,7 @@ 
Source code for brevitas.ops.autograd_ste_ops
'scalar_clamp_ste_impl',
     'tensor_clamp_ste_impl',
     'abs_binary_sign_grad_impl',
-    'dpu_round_ste_impl'
-]
+    'dpu_round_ste_impl']
 
 
 
[docs]class ScalarClampSteFn(Function):
@@ -486,7 +452,7 @@ 
Source code for brevitas.ops.autograd_ste_ops
[docs]class TensorClampSteFn(Function):
-    """ 
+    """
     Autograd function that implements :func:`~brevitas.function.ops.tensor_clamp` with a
     straight-through gradient estimator for the gradient of y w.r.t. to x, while the gradient of y
     w.r.t. to min_val and max_val is always None.
@@ -605,7 +571,7 @@ 
Source code for brevitas.ops.autograd_ste_ops
[docs]class CeilSteFn(Function):
-    """ 
+    """
     Autograd function that implements :func:`torch.ceil` with a straight-through gradient estimator.
 
     ``CeilSteFn.apply(*args)`` is first aliased to :func:`ceil_ste_impl(*args)
@@ -631,7 +597,7 @@ 
Source code for brevitas.ops.autograd_ste_ops
[docs]class FloorSteFn(Function):
-    """ 
+    """
     Autograd function that implements :func:`torch.floor` with a straight-through gradient estimator.
 
     ``FloorSteFn.apply(*args)`` is first aliased to :func:`floor_ste_impl(*args)
@@ -657,12 +623,12 @@ 
Source code for brevitas.ops.autograd_ste_ops
[docs]class BinarySignSteFn(Function):
-    """ 
+    """
     Autograd function that implements :func:`~brevitas.function.ops.binary_sign` with a
     straight-through gradient estimator.
 
-    ``BinarySignSteFn.apply(*args)`` is first aliased to 
-    :func:`binary_sign_ste_impl(*args)<brevitas.ops.autograd_ste_ops.binary_sign_ste_impl>` 
+    ``BinarySignSteFn.apply(*args)`` is first aliased to
+    :func:`binary_sign_ste_impl(*args)<brevitas.ops.autograd_ste_ops.binary_sign_ste_impl>`
     and then wrapped by :func:`~brevitas.function.ops_ste.binary_sign_ste` when env ``BREVITAS_JIT=0``.
     See :func:`~brevitas.function.ops_ste.binary_sign_ste` for details on the interface and
     examples.
@@ -689,7 +655,7 @@ 
Source code for brevitas.ops.autograd_ste_ops
[docs]class TernarySignSteFn(Function):
-    """ 
+    """
     Autograd function that implements :func:`torch.sign` with a straight-through gradient estimator.
 
     ``TernarySignSteFn.apply(*args)`` is first aliased to :func:`ternary_sign_ste_impl(*args)
diff --git a/docs/_modules/index.html b/docs/_modules/index.html
index f029f84df..09416cad5 100644
--- a/docs/_modules/index.html
+++ b/docs/_modules/index.html
@@ -5,7 +5,7 @@
   
     
     
-    Overview: module code — Brevitas 0.7.2.dev139+g0c2e90d documentation
+    Overview: module code — Brevitas 0.9.0 documentation
   
   
   
diff --git a/docs/_static/documentation_options.js b/docs/_static/documentation_options.js
index ffdf916a1..f725dd48e 100644
--- a/docs/_static/documentation_options.js
+++ b/docs/_static/documentation_options.js
@@ -1,6 +1,6 @@
 var DOCUMENTATION_OPTIONS = {
     URL_ROOT: document.getElementById("documentation_options").getAttribute('data-url_root'),
-    VERSION: '0.7.2.dev139+g0c2e90d',
+    VERSION: '0.9.0',
     LANGUAGE: 'en',
     COLLAPSE_INDEX: false,
     BUILDER: 'html',
diff --git a/docs/about.html b/docs/about.html
index aa4a7abd5..7bfdc1ce5 100644
--- a/docs/about.html
+++ b/docs/about.html
@@ -6,7 +6,7 @@
     
     
 
-    About — Brevitas 0.7.2.dev139+g0c2e90d documentation
+    About — Brevitas 0.9.0 documentation
   
   
   
diff --git a/docs/api_reference/brevitas.core.bit_width.html b/docs/api_reference/brevitas.core.bit_width.html
index 0bca61d16..dbc0ce72c 100644
--- a/docs/api_reference/brevitas.core.bit_width.html
+++ b/docs/api_reference/brevitas.core.bit_width.html
@@ -6,7 +6,7 @@
     
     
 
-    brevitas.core.bit_width package — Brevitas 0.7.2.dev139+g0c2e90d documentation
+    brevitas.core.bit_width package — Brevitas 0.9.0 documentation
   
   
   
@@ -386,7 +386,7 @@ 
Submodules
 

 class brevitas.core.bit_width.const.BitWidthConst(bit_width)[source]#

-
Bases: Module

+
Bases: Module

 
ScriptModule that returns a constant bit-width wrapped in a float torch.tensor.

 

 
Parameters:

@@ -422,7 +422,7 @@ 
Submodules
 
Return type:

-
Tensor

+
Tensor

 

 

 

@@ -434,10 +434,64 @@ 
Submodules
+

+class brevitas.core.bit_width.const.BitWidthStatefulConst(bit_width)[source]#

+
Bases: Module

+
ScriptModule that returns a constant bit-width wrapped in a float torch.tensor but retains the
+bit-width as part of the module state.

+

+
Parameters:

+
bit_width (int) – bit-width value.

+

+

+
Examples

+
>>> bit_width = BitWidthStatefulConst(8)
+>>> bit_width()
+tensor(8.)
+

+

+

+
Note

+
The BitWidthStatefulConst is a counterpart to BitWidthConst with the difference that the
+BitWidthStatefulConst retains the bit-width as part of the Module’s state. This means that it
+will be saved as part of a checkpoint.

+

+

+
Note

+
Maps to bit_width_impl_type == BitWidthImplType.STATEFUL_CONST == ‘STATEFUL_CONST’ ==
+‘stateful_const’ in higher-level APIs.

+

+

+

+forward()[source]#

+
Defines 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.

+

+

+
Return type:

+
Tensor

+

+

+

+
+

+

+training: bool#

+

+
+

+
 

 

 class brevitas.core.bit_width.const.MsbClampBitWidth(bit_width_to_remove_impl, min_overall_bit_width, max_overall_bit_width)[source]#

-
Bases: Module

+
Bases: Module

 

 

 forward(input_bit_width)[source]#

@@ -452,7 +506,7 @@ 
Submodules
 
Return type:

-
Tensor

+
Tensor

 

 

 

@@ -470,14 +524,14 @@ 
Submodules
 

 class brevitas.core.bit_width.parameter.BitWidthParameter(bit_width, min_bit_width=2, restrict_bit_width_impl=IntRestrictValue(   (float_to_int_impl): RoundSte() ), override_pretrained_bit_width=False)[source]#

-
Bases: Module

+
Bases: Module

 
ScriptModule that returns a learnable bit-width wrapped in a float torch.Tensor.

 

 
Parameters:

 
    
 
  • bit_width (int) – value to initialize the output learned bit-width.
    • 
 
  • min_bit_width (int) – lower bound for the output learned bit-width. Default: 2.
    • 
-
  • restrict_bit_width_impl (Module) – restrict the learned bit-width to a subset of values. Default: IntRestrictValue(RoundSte()).
    • 
+
  • restrict_bit_width_impl (Module) – restrict the learned bit-width to a subset of values. Default: IntRestrictValue(RoundSte()).
    • 
 
  • override_pretrained_bit_width (bool) – ignore pretrained bit-width loaded from a state dict. Default: False.
    • 
 

 

@@ -520,7 +574,7 @@ 
Submodules
 
Return type:

-
Tensor

+
Tensor

 

 

 

@@ -535,7 +589,7 @@ 
Submodules
 

 class brevitas.core.bit_width.parameter.RemoveBitwidthParameter(bit_width_to_remove, override_pretrained_bit_width=False, non_zero_epsilon=1e-06, remove_zero_bit_width=0.1)[source]#

-
Bases: Module

+
Bases: Module

 

 

 forward()[source]#

@@ -550,7 +604,7 @@ 
Submodules
 
Return type:

-
Tensor

+
Tensor

 

 

 

@@ -645,6 +699,35 @@ 
Submodules
+    
+     
+      
+       BitWidthStatefulConst
+      
+     
+    
+    
+   
    
  • 
     
      
diff --git a/docs/api_reference/brevitas.core.function_wrapper.html b/docs/api_reference/brevitas.core.function_wrapper.html
index 1f0bd93cf..2a5d6c940 100644
--- a/docs/api_reference/brevitas.core.function_wrapper.html
+++ b/docs/api_reference/brevitas.core.function_wrapper.html
@@ -6,7 +6,7 @@
     
     
 
-    brevitas.core.function_wrapper package — Brevitas 0.7.2.dev139+g0c2e90d documentation
+    brevitas.core.function_wrapper package — Brevitas 0.9.0 documentation
   
   
   
@@ -387,7 +387,7 @@ 
    Submodules
 
    
 class brevitas.core.function_wrapper.clamp.ClampMin(min_val)[source]#
    
-
    Bases: Module
    
+
    Bases: Module
    
 
    ScriptModule wrapper for clamp_min().
    
 
    Examples
    
 
    >>> clamp_min = ClampMin(min_val=-2.0)
@@ -419,8 +419,8 @@ 
    Submodules
 
    
 class brevitas.core.function_wrapper.clamp.ScalarClamp(min_val, max_val)[source]#
    
-
    Bases: Module
    
-
    ScriptModule wrapper for clamp().
    
+
    Bases: Module
    
+
    ScriptModule wrapper for clamp().
    
 
    Examples
    
 
    >>> scalar_clamp = ScalarClamp(min_val=-2.0, max_val=2.0)
 >>> scalar_clamp(torch.tensor([-3.0, 3.0]))
@@ -451,7 +451,7 @@ 
    Submodules
 
    
 class brevitas.core.function_wrapper.clamp.TensorClamp[source]#
    
-
    Bases: Module
    
+
    Bases: Module
    
 
    ScriptModule wrapper for tensor_clamp().
    
 
    Examples
    
 
    >>> tensor_clamp = TensorClamp()
@@ -489,7 +489,7 @@ 
    Submodules
 
    
 class brevitas.core.function_wrapper.misc.Identity[source]#
    
-
    Bases: Module
    
+
    Bases: Module
    
 
    Identity ScriptModule.
    
 
    Examples
    
 
    >>> identity = Identity()
@@ -513,7 +513,7 @@ 
    Submodules
 
    Return type:
    
-
    Tensor
    
+
    Tensor
    
 
    
 
 
    
@@ -528,7 +528,7 @@ 
    Submodules
 
    
 class brevitas.core.function_wrapper.misc.InplaceLogTwo[source]#
    
-
    Bases: Module
    
+
    Bases: Module
    
 
    Module wrapper for log2_().
    
 
    Examples
    
 
    >>> inplace_log_two = InplaceLogTwo()
@@ -554,7 +554,7 @@ 
    Submodules
 
    Return type:
    
-
    Tensor
    
+
    Tensor
    
 
    
 
 
    
@@ -569,8 +569,8 @@ 
    Submodules
 
    
 class brevitas.core.function_wrapper.misc.LogTwo[source]#
    
-
    Bases: Module
    
-
    ScriptModule wrapper for log2().
    
+
    Bases: Module
    
+
    ScriptModule wrapper for log2().
    
 
    Examples
    
 
    >>> log_two = LogTwo()
 >>> x = torch.tensor(8.0)
@@ -592,7 +592,7 @@ 
    Submodules
 
    Return type:
    
-
    Tensor
    
+
    Tensor
    
 
    
 
 
    
@@ -607,7 +607,7 @@ 
    Submodules
 
    
 class brevitas.core.function_wrapper.misc.PowerOfTwo[source]#
    
-
    Bases: Module
    
+
    Bases: Module
    
 
    ScriptModule implementation of 2.0 ** x.
    
 
    Examples
    
 
    >>> power_of_two = PowerOfTwo()
@@ -630,7 +630,7 @@ 
    Submodules
 
    Return type:
    
-
    Tensor
    
+
    Tensor
    
 
    
 
 
    
@@ -649,7 +649,7 @@ 
    Submodules
 
    
 class brevitas.core.function_wrapper.ops_ste.CeilSte[source]#
    
-
    Bases: Module
    
+
    Bases: Module
    
 
    ScriptModule wrapper for ceil_ste().
    
 
    
 
    
@@ -675,7 +675,7 @@ 
    Submodules
 
    
 class brevitas.core.function_wrapper.ops_ste.DPURoundSte[source]#
    
-
    Bases: Module
    
+
    Bases: Module
    
 
    ScriptModule wrapper for dpu_round_ste().
    
 
    
 
    
@@ -701,7 +701,7 @@ 
    Submodules
 
    
 class brevitas.core.function_wrapper.ops_ste.FloorSte[source]#
    
-
    Bases: Module
    
+
    Bases: Module
    
 
    ScriptModule wrapper for floor_ste().
    
 
    
 
    
@@ -727,7 +727,7 @@ 
    Submodules
 
    
 class brevitas.core.function_wrapper.ops_ste.InplaceTensorClampSte[source]#
    
-
    Bases: Module
    
+
    Bases: Module
    
 
    ScriptModule wrapper for tensor_clamp_ste_().
    
 
    
 
    
@@ -753,7 +753,7 @@ 
    Submodules
 
    
 class brevitas.core.function_wrapper.ops_ste.RoundSte[source]#
    
-
    Bases: Module
    
+
    Bases: Module
    
 
    ScriptModule wrapper for round_ste().
    
 
    
 
    
@@ -779,7 +779,7 @@ 
    Submodules
 
    
 class brevitas.core.function_wrapper.ops_ste.RoundToZeroSte[source]#
    
-
    Bases: Module
    
+
    Bases: Module
    
 
    ScriptModule wrapper for round_to_zero_ste().
    
 
    
 
    
@@ -805,7 +805,7 @@ 
    Submodules
 
    
 class brevitas.core.function_wrapper.ops_ste.ScalarClampMinSte(min_val)[source]#
    
-
    Bases: Module
    
+
    Bases: Module
    
 
    ScriptModule wrapper for scalar_clamp_min_ste().
    
 
    
 
    
@@ -831,7 +831,7 @@ 
    Submodules
 
    
 class brevitas.core.function_wrapper.ops_ste.TensorClampSte[source]#
    
-
    Bases: Module
    
+
    Bases: Module
    
 
    ScriptModule wrapper for tensor_clamp_ste().
    
 
    
 
    
@@ -861,7 +861,7 @@ 
    Submodules
 
    
 class brevitas.core.function_wrapper.shape.OverBatchOverOutputChannelView[source]#
    
-
    Bases: Module
    
+
    Bases: Module
    
 
    ScriptModule to compute the over_batch_over_output_channels()
 view of an input tensor.
    
 
    Examples
    
@@ -895,7 +895,7 @@ 
    Submodules
 
    
 class brevitas.core.function_wrapper.shape.OverBatchOverTensorView[source]#
    
-
    Bases: Module
    
+
    Bases: Module
    
 
    ScriptMoodule to compute the over_batch_over_tensor() view of an
 input tensor.
    
 
    Examples
    
@@ -929,7 +929,7 @@ 
    Submodules
 
    
 class brevitas.core.function_wrapper.shape.OverOutputChannelView(permute_dims)[source]#
    
-
    Bases: Module
    
+
    Bases: Module
    
 
    ScriptMoodule to compute the over_output_channels() view of an
 input tensor.
    
 
    Examples
    
@@ -963,7 +963,7 @@ 
    Submodules
 
    
 class brevitas.core.function_wrapper.shape.OverTensorView[source]#
    
-
    Bases: Module
    
+
    Bases: Module
    
 
    ScriptMoodule to compute the over_tensor() view of an input tensor.
    
 
    Examples
    
 
    >>> view_module = OverTensorView()
@@ -996,7 +996,7 @@ 
    Submodules
 
    
 class brevitas.core.function_wrapper.shape.PermuteDims(permute_dims)[source]#
    
-
    Bases: Module
    
+
    Bases: Module
    
 
    
 
    
 forward(x)[source]#
    
diff --git a/docs/api_reference/brevitas.core.html b/docs/api_reference/brevitas.core.html
index 7587a10d9..59d5e924a 100644
--- a/docs/api_reference/brevitas.core.html
+++ b/docs/api_reference/brevitas.core.html
@@ -6,7 +6,7 @@
     
     
 
-    brevitas.core package — Brevitas 0.7.2.dev139+g0c2e90d documentation
+    brevitas.core package — Brevitas 0.9.0 documentation
   
   
   
@@ -390,6 +390,11 @@ 
    SubpackagesBitWidthConst.training
    • 
 
 
+
  • BitWidthStatefulConst
+
    • 
 
  • MsbClampBitWidth
 
    • 
+
  • DecoupledRescalingIntQuantWithInput
+
    • 
 
  • PrescaledRestrictIntQuant
 
    • 
+
  • L1Norm
+
    • 
+
  • L2Norm
+
    • 
 
  • MeanLearnedSigmaStd