Skip to main content
`Ctrl`+`K`

# Source code for brevitas.core.stats.stats_op

    # Copyright (C) 2023, Advanced Micro Devices, Inc. All rights reserved.
    # SPDX-License-Identifier: BSD-3-Clause

    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.function_wrapper.misc import Identity
    from brevitas.core.function_wrapper.ops_ste import ScalarClampMinSte
    from brevitas.core.utils import StatelessBuffer
    from brevitas.function.ops import max_int
    from brevitas.quant_tensor import _unpack_quant_tensor
    # Use custom implementation of kthvalue as work around to (b)float16 kernel limitations
    from brevitas.utils.torch_utils import kthvalue

    from .stats_wrapper import SCALAR_SHAPE

    DEFAULT_STD_DEV_EPSILON = 1e-8

    class NegativeMinOrZero(brevitas.jit.ScriptModule):
        __constants__ = ['stats_reduce_dim', 'keepdim']

        def __init__(
                self,
                stats_reduce_dim: Optional[int] = None,
                dtype: Optional[torch.dtype] = None,
                device: Optional[torch.device] = None,
                keepdim: bool = False) -> None:
            super(NegativeMinOrZero, self).__init__()
            self.stats_reduce_dim = stats_reduce_dim
            self.zero = StatelessBuffer(torch.tensor(0.0, dtype=dtype, device=device))
            self.keepdim = keepdim

        @brevitas.jit.script_method
        def forward(self, x: Tensor) -> Tensor:
            if self.stats_reduce_dim is None:
                min_val = torch.min(x)
            else:
                min_val = torch.min(x, dim=self.stats_reduce_dim, keepdim=self.keepdim)[0]
            min_val = torch.clamp(min_val, max=self.zero())
            return min_val

    class AbsPercentile(brevitas.jit.ScriptModule):
        __constants__ = ['q', 'stats_reduce_dim', 'keepdim']

        def __init__(
                self,
                high_percentile_q: float,
                stats_reduce_dim: Optional[int],
                percentile_q=None,
                keepdim: bool = False):
            super(AbsPercentile, self).__init__()
            if percentile_q is not None:
                raise RuntimeError("percentile_q is deprecated, please pass high_percentile_q.")
            assert high_percentile_q <= 100, "q has to be a percentage"
            self.q = high_percentile_q
            self.stats_reduce_dim = stats_reduce_dim
            self.keepdim = keepdim

        @brevitas.jit.script_method
        def forward(self, x: Tensor):
            if self.stats_reduce_dim is None:
                # k is 1-indexed, so round away from zero
                k = int(math.floor(.01 * self.q * x.numel() + 0.5))
                result = kthvalue(x.abs().view(-1), k)[0]
            else:
                # assuming x is two dimensional, get the other dimension
                assert len(x.size()) == 2, "Only 2-dim input is supported."
                other_dim = abs(self.stats_reduce_dim - 1)
                dim_slice = torch.narrow(x, dim=other_dim, start=0, length=1)
                # k is 1-indexed, so round away from zero
                k = int(math.floor(.01 * self.q * dim_slice.numel() + 0.5))
                result = kthvalue(x.abs(), k, dim=self.stats_reduce_dim, keepdim=self.keepdim)[0]
            return result

    class NegativePercentileOrZero(brevitas.jit.ScriptModule):
        __constants__ = ['stats_reduce_dim', 'q', 'keepdim']

        def __init__(
                self,
                low_percentile_q,
                stats_reduce_dim: Optional[int] = None,
                dtype: Optional[torch.dtype] = None,
                device: Optional[torch.device] = None,
                keepdim: bool = False) -> None:
            super(NegativePercentileOrZero, self).__init__()
            self.stats_reduce_dim = stats_reduce_dim
            self.q = low_percentile_q
            self.zero = StatelessBuffer(torch.tensor(0.0, dtype=dtype, device=device))
            self.keepdim = keepdim

        @brevitas.jit.script_method
        def forward(self, x: Tensor) -> Tensor:
            if self.stats_reduce_dim is None:
                # k is 1-indexed, so round away from zero
                k = int(math.ceil(.01 * self.q * x.numel()))
                result = kthvalue(x.view(-1), k)[0]
            else:
                # assuming x is two dimensional, get the other dimension
                assert len(x.size()) == 2, "Only 2-dim input is supported."
                other_dim = abs(self.stats_reduce_dim - 1)
                dim_slice = torch.narrow(x, dim=other_dim, start=0, length=1)
                # k is 1-indexed, so round away from zero
                k = int(math.ceil(.01 * self.q * dim_slice.numel()))
                result = kthvalue(x, k, dim=self.stats_reduce_dim, keepdim=self.keepdim)[0]
            result = torch.clamp(result, max=self.zero())
            return result

    class PercentileInterval(brevitas.jit.ScriptModule):
        __constants__ = ['stats_reduce_dim', 'low_q', 'high_q', 'keepdim']

        def __init__(
                self,
                low_percentile_q,
                high_percentile_q,
                stats_reduce_dim: Optional[int] = None,
                keepdim: bool = False,
                dtype: Optional[torch.dtype] = None,
                device: Optional[torch.device] = None) -> None:
            super(PercentileInterval, self).__init__()
            self.stats_reduce_dim = stats_reduce_dim
            self.low_q = low_percentile_q
            self.high_q = high_percentile_q
            self.keepdim = keepdim
            self.zero = StatelessBuffer(torch.tensor(0.0, dtype=dtype, device=device))

        @brevitas.jit.script_method
        def forward(self, x: Tensor) -> Tensor:
            if self.stats_reduce_dim is None:
                low_k = int(math.ceil(.01 * self.low_q * x.numel()))
                # k is 1-indexed, so round away from zero
                high_k = int(math.floor(.01 * self.high_q * x.numel() + 0.5))
                low_result = kthvalue(x.view(-1), low_k)[0]
                high_result = kthvalue(x.view(-1), high_k)[0]
            else:
                # assuming x is two dimensional, get the other dimension
                assert len(x.size()) == 2, "Only 2-dim input is supported."
                other_dim = abs(self.stats_reduce_dim - 1)
                dim_slice = torch.narrow(x, dim=other_dim, start=0, length=1)
                low_k = int(math.ceil(.01 * self.low_q * dim_slice.numel()))
                # k is 1-indexed, so round away from zero
                high_k = int(math.floor(.01 * self.high_q * dim_slice.numel() + 0.5))
                low_result = kthvalue(x, low_k, dim=self.stats_reduce_dim, keepdim=self.keepdim)[0]
                high_result = kthvalue(x, high_k, dim=self.stats_reduce_dim, keepdim=self.keepdim)[0]
            # We need to make sure the lower bound is not positive to align with zero-point statistics
            low_result = torch.clamp(low_result, max=self.zero())
            interval = high_result - low_result
            abs_interval = torch.abs(interval)
            return abs_interval

    class AbsMax(brevitas.jit.ScriptModule):
        __constants__ = ['stats_reduce_dim']

        def __init__(self, stats_reduce_dim: Optional[int] = None, keepdim: bool = False) -> None:
            super(AbsMax, self).__init__()
            self.stats_reduce_dim = stats_reduce_dim
            self.keepdim = keepdim

        @brevitas.jit.script_method
        def forward(self, x: Tensor):
            if self.stats_reduce_dim is None:
                return torch.max(torch.abs(x))
            else:
                return torch.max(torch.abs(x), dim=self.stats_reduce_dim, keepdim=self.keepdim)[0]

    class AbsMinMax(brevitas.jit.ScriptModule):
        __constants__ = ['stats_reduce_dim', 'keepdim']

        def __init__(
                self,
                stats_reduce_dim: Optional[int] = None,
                keepdim: bool = False,
                dtype: Optional[torch.dtype] = None,
                device: Optional[torch.device] = None) -> None:
            super(AbsMinMax, self).__init__()
            self.stats_reduce_dim = stats_reduce_dim
            self.keepdim = keepdim
            self.zero = StatelessBuffer(torch.tensor(0.0, dtype=dtype, device=device))

        @brevitas.jit.script_method
        def forward(self, x: Tensor):
            if self.stats_reduce_dim is None:
                max_val = torch.max(x)
                min_val = torch.min(x)
            else:
                max_val = torch.max(x, dim=self.stats_reduce_dim, keepdim=self.keepdim)[0]
                min_val = torch.min(x, dim=self.stats_reduce_dim, keepdim=self.keepdim)[0]
            # We need to make sure the lower bound is not positive to align with zero-point statistics
            min_val = torch.clamp(min_val, max=self.zero())
            return torch.abs(max_val - min_val)

    class AbsMaxAve(brevitas.jit.ScriptModule):
        __constants__ = ['stats_reduce_dim']

        def __init__(self, stats_reduce_dim: int) -> None:
            super(AbsMaxAve, self).__init__()
            self.stats_reduce_dim = stats_reduce_dim

        @brevitas.jit.script_method
        def forward(self, x: Tensor):
            return torch.mean(torch.max(torch.abs(x), dim=self.stats_reduce_dim)[0])

    class AbsMaxL2(brevitas.jit.ScriptModule):
        __constants__ = ['stats_reduce_dim']

        def __init__(self, stats_reduce_dim: int) -> None:
            super(AbsMaxL2, self).__init__()
            self.stats_reduce_dim = stats_reduce_dim

        @brevitas.jit.script_method
        def forward(self, x: torch.Tensor):
            per_channel_max = torch.max(torch.abs(x), dim=self.stats_reduce_dim)[0]
            out = torch.norm(per_channel_max, p=2)
            out = out / math.sqrt(per_channel_max.view(-1).shape[0])
            return out

    class AbsAve(brevitas.jit.ScriptModule):
        __constants__ = ['stats_reduce_dim']

        def __init__(self, stats_reduce_dim: Optional[int] = None) -> None:
            super(AbsAve, self).__init__()
            self.stats_reduce_dim = stats_reduce_dim

        @brevitas.jit.script_method
        def forward(self, x: Tensor):
            if self.stats_reduce_dim is None:
                return torch.mean(torch.abs(x))
            else:
                return torch.mean(torch.abs(x), dim=self.stats_reduce_dim)

    class MeanSigmaStd(brevitas.jit.ScriptModule):

        def __init__(
                self,
                sigma: float,
                stats_reduce_dim: Optional[int] = None,
                std_dev_epsilon: float = DEFAULT_STD_DEV_EPSILON,
                dtype: Optional[torch.dtype] = None,
                device: Optional[torch.device] = None) -> None:
            super(MeanSigmaStd, self).__init__()
            self.impl = _MeanSigmaStdImpl(stats_reduce_dim, std_dev_epsilon)
            self.sigma = StatelessBuffer(torch.tensor(sigma, dtype=dtype, device=device))

        @brevitas.jit.script_method
        def forward(self, x: Tensor):
            sigma = self.sigma()
            out = self.impl(x, sigma)
            return out

    class _MeanSigmaStdImpl(brevitas.jit.ScriptModule):
        __constants__ = ['stats_reduce_dim', 'output_shape', 'epsilon']

        def __init__(
                self,
                stats_reduce_dim: Optional[int] = None,
                std_dev_epsilon: float = DEFAULT_STD_DEV_EPSILON) -> None:
            super(_MeanSigmaStdImpl, self).__init__()
            self.stats_reduce_dim = stats_reduce_dim
            self.epsilon = std_dev_epsilon

        @brevitas.jit.script_method
        def forward(self, x: Tensor, sigma: Tensor):
            abs_val = torch.abs(x)
            if self.stats_reduce_dim is None:
                mean_val = torch.mean(abs_val)
                std_val = torch.sqrt(torch.var(abs_val) + self.epsilon)
            else:
                mean_val = torch.mean(torch.abs(x), dim=self.stats_reduce_dim)
                std_val = torch.sqrt(torch.var(abs_val, dim=self.stats_reduce_dim) + self.epsilon)
                mean_val = mean_val.view(-1)
                std_val = std_val.view(-1)
            return mean_val + sigma * std_val

    class MeanLearnedSigmaStd(brevitas.jit.ScriptModule):

        def __init__(
                self,
                sigma: float,
                stats_output_shape: Tuple[int, ...],
                stats_reduce_dim: Optional[int] = None,
                std_dev_epsilon: float = DEFAULT_STD_DEV_EPSILON,
                dtype: Optional[torch.dtype] = None,
                device: Optional[torch.device] = None) -> None:
            super(MeanLearnedSigmaStd, self).__init__()
            self.impl = _MeanSigmaStdImpl(stats_reduce_dim, std_dev_epsilon)
            if stats_output_shape == SCALAR_SHAPE:
                self.value = Parameter(torch.tensor(sigma, dtype=dtype, device=device))
            else:
                self.value = Parameter(
                    torch.full(stats_output_shape, sigma, dtype=dtype, device=device))

        @brevitas.jit.script_method
        def forward(self, x: Tensor):
            sigma = self.sigma.view(self.sigma.shape)  # trick to get a tensor type
            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):
            value_key = prefix + 'sigma'
            retrocomp_value_key = prefix + 'learned_sigma'
            if retrocomp_value_key in state_dict:  # retrocompatibility
                state_dict[value_key] = state_dict.pop(retrocomp_value_key)
            super(MeanLearnedSigmaStd, self)._load_from_state_dict(
                state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
            sigma_key = prefix + 'sigma'
            if config.IGNORE_MISSING_KEYS and sigma_key in missing_keys:
                missing_keys.remove(sigma_key)

    class KLMinimizerThreshold(torch.nn.Module):
        """
        Based on:
        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):
            super(KLMinimizerThreshold, self).__init__()
            self.num_bins = num_bins
            self.smoothing_eps = smoothing_eps
            self.signed = signed
            self.bit_width_impl = bit_width_impl
            self.absmax_impl = AbsMax()

        def smooth_normalize_distribution(self, p, eps):
            is_zeros = (p == 0).float()
            n_zeros = is_zeros.sum()
            n_nonzeros = torch.numel(p) - n_zeros
            if not n_nonzeros:
                return None
            eps1 = eps * n_zeros / n_nonzeros
            hist = p.float()
            hist += eps * is_zeros + (-eps1) * n_nonzeros
            dist = torch.distributions.categorical.Categorical(logits=hist)
            return dist

        def forward(self, x: Tensor):
            absmax = self.absmax_impl(x)
            bit_width = self.bit_width_impl()
            num_quantized_bins = max_int(self.signed, False, bit_width).int()
            thresholds = torch.zeros(self.num_bins // 2 + 1 - num_quantized_bins // 2, device=x.device)
            divergence = torch.zeros_like(thresholds)
            quantized_bins = torch.zeros(num_quantized_bins, device=x.device)
            hist = torch.histc(x, bins=self.num_bins, min=-absmax, max=absmax).int()
            hist_edges = torch.linspace(-absmax, absmax, self.num_bins + 1)
            for i in range(num_quantized_bins // 2, self.num_bins // 2 + 1):
                p_bin_idx_start = self.num_bins // 2 - i
                p_bin_idx_stop = self.num_bins // 2 + i + 1
                thresholds[i - num_quantized_bins // 2] = hist_edges[p_bin_idx_stop]
                sliced_nd_hist = hist[p_bin_idx_start:p_bin_idx_stop]
                p = sliced_nd_hist.clone()
                left_outlier_count = torch.sum(hist[0:p_bin_idx_start])
                p[0] += left_outlier_count
                right_outlier_count = torch.sum(hist[p_bin_idx_stop:])
                p[-1] += right_outlier_count
                is_nonzeros = (sliced_nd_hist != 0).float()
                num_merged_bins = torch.numel(p) // num_quantized_bins
                for j in range(num_quantized_bins):
                    start = j * num_merged_bins
                    stop = start + num_merged_bins
                    quantized_bins[j] = sliced_nd_hist[start:stop].sum()
                quantized_bins[-1] += sliced_nd_hist[num_quantized_bins * num_merged_bins:].sum()
                q = torch.zeros_like(p, dtype=torch.float32, device=x.device)
                for j in range(num_quantized_bins):
                    start = j * num_merged_bins
                    if j == num_quantized_bins - 1:
                        stop = -1
                    else:
                        stop = start + num_merged_bins
                    norm = is_nonzeros[start:stop].sum()
                    if norm != 0:
                        q[start:stop] = quantized_bins[j] / norm
                q[sliced_nd_hist == 0] = 0.
                p = self.smooth_normalize_distribution(p, self.smoothing_eps)
                q = self.smooth_normalize_distribution(q, self.smoothing_eps)
                if q is None:
                    divergence[i - num_quantized_bins // 2] = float('inf')
                else:
                    divergence[i - num_quantized_bins // 2] = torch.distributions.kl.kl_divergence(p, q)
            min_divergence_idx = torch.argmin(divergence)
            opt_threshold = thresholds[min_divergence_idx]
            return opt_threshold

    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

        @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)

    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

        @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)

    def _set_local_loss_mode(module, enabled):
        for m in module.modules():
            if hasattr(m, 'local_loss_mode'):
                m.local_loss_mode = enabled

    def _set_observer_mode(module, enabled, previous_observer_mode):
        for m in module.modules():
            if hasattr(m, 'observer_only'):
                previous_observer_mode[m] = m.observer_only
                m.observer_only = enabled

    def _restore_observer_mode(module, previous_observer_mode):
        for m in module.modules():
            if hasattr(m, 'observer_only'):
                m.observer_only = previous_observer_mode[m]

    class MSE(torch.nn.Module):
        # References:
        # https://github.com/cornell-zhang/dnn-quant-ocs/blob/master/distiller/quantization/clip.py
        # https://github.com/wimh966/outlier_suppression/blob/main/quant_transformer/quantization/observer.py

        def __init__(
                self,
                proxy_module,
                mse_init_op,
                inner_stats_input_view_shape_impl: torch.nn.Module,
                stats_reduce_dim: Optional[int] = None,
                mse_search_method: str = 'fibonacci',
                mse_iters: int = 20):
            super(MSE, self).__init__()
            self.mse_init_op = mse_init_op
            self.input_view_shape_impl = inner_stats_input_view_shape_impl
            self.proxy_forward = proxy_module.forward
            self.previous_observer_mode = dict()
            self.set_local_loss_mode = lambda enabled: _set_local_loss_mode(proxy_module, enabled)
            self.set_observer_mode = lambda enabled: _set_observer_mode(
                proxy_module, enabled, self.previous_observer_mode)
            self.restore_observer_mode = lambda: _restore_observer_mode(
                proxy_module, self.previous_observer_mode)
            self.internal_candidate = None
            self.num = mse_iters
            self.search_method = mse_search_method
            self.stats_reduce_dim = stats_reduce_dim
            self.local_loss_mode: bool = False

        def mse_loss_fn(self, x, quant_value):
            loss = torch.nn.functional.mse_loss(x, quant_value, reduction='none')
            if self.stats_reduce_dim is not None:
                # stats_reduce_dim applies to the permuted and reshaped tensor
                loss = self.input_view_shape_impl(loss)
                loss = torch.sum(loss, dim=self.stats_reduce_dim)
            else:
                loss = torch.sum(loss)
            return loss

        def evaluate_loss(self, x, candidate):
            self.internal_candidate = candidate
            # Set to local_loss_mode before calling the proxy
            self.set_local_loss_mode(True)
            self.set_observer_mode(False)
            quant_value = self.proxy_forward(x)
            quant_value = _unpack_quant_tensor(quant_value)
            loss = self.mse_loss_fn(x, quant_value)
            self.set_local_loss_mode(False)
            self.restore_observer_mode()
            return loss

        def mse_grid_search(self, xl, x):
            best_loss = torch.tensor(float('inf'), device=x.device, dtype=x.dtype)
            best_candidate = xl
            for i in range(2, self.num + 1):
                candidate = (xl * i).detach()
                loss = self.evaluate_loss(x, candidate)
                best_candidate = torch.where(loss < best_loss, candidate, best_candidate)
                best_loss = torch.min(loss, best_loss)
            return best_candidate

        def mse_fib_search(self, xl, xr, x):

            def fib_seq(n):
                if n <= 0:
                    return [0]
                seq = [0, 1]
                while len(seq) <= n:
                    next = seq[-1] + seq[-2]
                    seq.append(next)
                return seq

            # vectorized variant of
            # https://indrag49.github.io/Numerical-Optimization/solving-one-dimensional-optimization-problems.html#fibonacci-search-method
            F = fib_seq(self.num)
            L0 = xr - xl
            Li = (F[self.num - 2] / F[self.num]) * L0
            for i in range(2, self.num + 1):
                x1 = torch.where(Li > L0 / 2, xr - Li, xl + Li)
                x2 = torch.where(Li > L0 / 2, xl + Li, xr - Li)
                f1, f2 = self.evaluate_loss(x, x1), self.evaluate_loss(x, x2)
                xr = torch.where(f1 <= f2, x2, xr)
                xl = torch.where(f1 >= f2, x1, xl)
                Li = (F[self.num - i] / F[self.num - (i - 2)]) * torch.where(f1 != f2, L0, xr - xl)
                L0 = xr - xl
            return torch.where(f1 <= f2, x1, x2)

        def mse_search(self, x):
            x_view = self.input_view_shape_impl(x)
            init = self.mse_init_op(x_view).detach()
            base = init / self.num
            if self.search_method == 'grid':
                best_candidate = self.mse_grid_search(base, x)
            elif self.search_method == 'fibonacci':
                best_candidate = self.mse_fib_search(base, init, x)
            else:
                raise ValueError(f"Search method {self.search_method} not supported.")
            # Save for evaluation by other modules (e.g. zp) invoking local loss mode
            self.internal_candidate = best_candidate
            return best_candidate

        def forward(self, x):
            if not self.local_loss_mode:
                with torch.no_grad():
                    return self.mse_search(x)
            else:
                # This is invoked for the zero-point whenever scale is being optimized first
                if self.internal_candidate is None:
                    x = self.input_view_shape_impl(x)
                    self.internal_candidate = self.mse_init_op(x).detach()
                return self.internal_candidate

    class HalfQuadraticOptimizerScale(torch.nn.Module):
        # References:
        # https://mobiusml.github.io/hqq_blog/
        # https://github.com/mobiusml/hqq?tab=readme-ov-file

        def __init__(
                self,
                proxy_module,
                hqo_init_op_scale,
                keepdim: bool,
                inner_stats_input_view_shape_impl: torch.nn.Module,
                scaling_min_val: Optional[float] = None,
                stats_reduce_dim: Optional[int] = None,
                int_scaling_impl=None,
                bit_width_impl=None,
                hqo_beta_scale: float = 1e5,
                hqo_kappa_scale: float = 1.01,
                hqo_lp_norm_scale: float = .7,
                hqo_iters_scale: int = 1000):
            super(HalfQuadraticOptimizerScale, self).__init__()
            self.hqo_init_op = hqo_init_op_scale
            self.input_view_shape_impl = inner_stats_input_view_shape_impl
            self.proxy_forward = proxy_module.forward
            self.previous_observer_mode = dict()
            self.set_local_loss_mode = lambda enabled: _set_local_loss_mode(proxy_module, enabled)
            self.set_observer_mode = lambda enabled: _set_observer_mode(
                proxy_module, enabled, self.previous_observer_mode)
            self.restore_observer_mode = lambda: _restore_observer_mode(
                proxy_module, self.previous_observer_mode)
            self.internal_candidate = None
            self.hqo_iters = hqo_iters_scale
            self.stats_reduce_dim = stats_reduce_dim
            self.local_loss_mode: bool = False

            self.beta = hqo_beta_scale
            self.kappa = hqo_kappa_scale
            self.lp_norm = hqo_lp_norm_scale

            self.int_scaling_impl = int_scaling_impl
            self.msb_clamp_bit_width_impl = bit_width_impl
            if scaling_min_val is not None and scaling_min_val != 0:
                self.clamp_min_ste = ScalarClampMinSte(scaling_min_val)
            else:
                self.clamp_min_ste = Identity()
            self.keepdim = keepdim

        def parameter_search(self, xl, x):
            best_loss = torch.tensor(float('inf'), device=x.device, dtype=x.dtype)
            candidate = xl
            best_candidate = candidate
            beta = self.beta
            with torch.no_grad():
                for i in range(0, self.hqo_iters):
                    self.internal_candidate = candidate
                    self.set_local_loss_mode(True)
                    self.set_observer_mode(False)
                    quant_tensor = self.proxy_forward(x).detach()
                    self.set_local_loss_mode(False)
                    self.restore_observer_mode()
                    loss = torch.abs(quant_tensor.value - x).mean()

                    best_candidate = torch.where(loss < best_loss, candidate, best_candidate)
                    if loss >= best_loss:
                        break
                    best_loss = torch.min(loss, best_loss)
                    W_e = shrink_lp_op(x - quant_tensor.value, beta, self.lp_norm)
                    zero_point = quant_tensor.zero_point
                    num = self.input_view_shape_impl(x - W_e).detach()
                    den = self.input_view_shape_impl(
                        torch.round(quant_tensor.value / quant_tensor.scale) - zero_point).detach()
                    mask = (num != 0.) & (den != 0.)
                    if self.stats_reduce_dim is None:
                        candidate = masked_median(num / den, mask)
                    else:
                        candidate = masked_median(
                            num / den, mask, dim=self.stats_reduce_dim, keepdim=self.keepdim)
                    candidate = candidate.type_as(self.internal_candidate)
                    candidate = self.clamp_min_ste(candidate)
                    bit_width = self.msb_clamp_bit_width_impl()
                    int_threshold = self.int_scaling_impl(bit_width)
                    candidate = candidate * int_threshold
                    candidate[torch.isnan(candidate)] = self.internal_candidate[torch.isnan(candidate)]
                    candidate[torch.isinf(candidate)] = self.internal_candidate[torch.isinf(candidate)]
                    beta *= self.kappa
            return best_candidate

        def optimize(self, x):
            x_view = self.input_view_shape_impl(x)

            init = self.hqo_init_op(x_view).detach()
            best_candidate = self.parameter_search(init, x_view)

            # Save for evaluation by other modules (e.g. zp) invoking local loss mode
            self.internal_candidate = best_candidate.detach()
            torch.cuda.empty_cache()
            return best_candidate

        def forward(self, x):
            if not self.local_loss_mode:
                with torch.no_grad():
                    return self.optimize(x)
            else:
                # This is invoked for the zero-point whenever scale is being optimized first
                if self.internal_candidate is None:
                    x = self.input_view_shape_impl(x)
                    self.internal_candidate = self.hqo_init_op(x).detach()
                return self.internal_candidate

    class HalfQuadraticOptimizerZeroPoint(torch.nn.Module):
        # References:
        # https://mobiusml.github.io/hqq_blog/
        # https://github.com/mobiusml/hqq?tab=readme-ov-file

        def __init__(
                self,
                proxy_module,
                keepdim: bool,
                hqo_init_op_zp: torch.nn.Module,
                inner_stats_input_view_shape_impl: torch.nn.Module,
                stats_reduce_dim: Optional[int] = None,
                hqo_beta_zp: float = 1e0,
                hqo_kappa_zp: float = 1.01,
                hqo_lp_norm_zp: float = .5,
                hqo_iters_zp: int = 1000):
            super(HalfQuadraticOptimizerZeroPoint, self).__init__()
            self.hqo_init_op_zp = hqo_init_op_zp
            self.input_view_shape_impl = inner_stats_input_view_shape_impl
            self.proxy_forward = proxy_module.forward
            self.previous_observer_mode = dict()
            self.set_local_loss_mode = lambda enabled: _set_local_loss_mode(proxy_module, enabled)
            self.set_observer_mode = lambda enabled: _set_observer_mode(
                proxy_module, enabled, self.previous_observer_mode)
            self.restore_observer_mode = lambda: _restore_observer_mode(
                proxy_module, self.previous_observer_mode)
            self.internal_candidate = None
            self.stats_reduce_dim = stats_reduce_dim
            self.local_loss_mode: bool = False
            self.beta = hqo_beta_zp
            self.kappa = hqo_kappa_zp
            self.lp_norm = hqo_lp_norm_zp
            self.hqo_iters = hqo_iters_zp
            self.keepdim = keepdim

        def parameter_search(self, xl, x):
            best_loss = torch.tensor(float('inf'), device=x.device, dtype=x.dtype)
            candidate = xl
            best_candidate = candidate
            with torch.no_grad():
                for i in range(0, self.hqo_iters):
                    self.internal_candidate = candidate
                    self.set_local_loss_mode(True)
                    self.set_observer_mode(False)
                    quant_tensor = self.proxy_forward(x).detach()
                    self.set_local_loss_mode(False)
                    self.restore_observer_mode()
                    qt_value = self.input_view_shape_impl(quant_tensor.value)
                    qt_scale = self.input_view_shape_impl(quant_tensor.scale)
                    qt_zp = self.input_view_shape_impl(quant_tensor.zero_point)
                    qt_int = qt_value / qt_scale + qt_zp
                    loss = torch.abs(qt_value - x).mean()
                    best_candidate = torch.where(loss < best_loss, candidate, best_candidate)
                    if loss >= best_loss:
                        break
                    best_loss = torch.min(loss, best_loss)
                    W_e = shrink_lp_op(x - qt_value, self.beta, self.lp_norm)

                    # Compared to the original formulation, the value we're looking for is:
                    # - scaled by qt_scale
                    # - opposite sign
                    val = self.input_view_shape_impl((x - W_e) - qt_int * qt_scale)

                    if self.stats_reduce_dim is None:
                        candidate = torch.mean(val)
                    else:
                        candidate = torch.mean(val, dim=self.stats_reduce_dim, keepdim=self.keepdim)
                    self.beta *= self.kappa
            return best_candidate

        def optimize(self, x):
            x_view = self.input_view_shape_impl(x)

            init = self.hqo_init_op_zp(x_view).detach()

            best_candidate = self.parameter_search(init, x)

            # Save for evaluation by other modules (e.g. zp) invoking local loss mode
            self.internal_candidate = best_candidate.detach()
            torch.cuda.empty_cache()
            return best_candidate

        def forward(self, x):
            if not self.local_loss_mode:
                with torch.no_grad():
                    return self.optimize(x)
            else:
                # This is invoked for the zero-point whenever scale is being optimized first
                if self.internal_candidate is None:
                    x = self.input_view_shape_impl(x)
                    self.internal_candidate = self.hqo_init_op_zp(x).detach()
                return self.internal_candidate

    def masked_median(x, mask, dim=None, keepdim=False):
        """Compute the median of tensor x along dim, ignoring values where mask is False.
        x and mask need to be broadcastable.

        Args:
            x (Tensor): Tensor to compute median of.
            mask (BoolTensor): Same shape as x with True where x is valid and False
                where x should be masked. Mask should not be all False in any column of
                dimension dim to avoid NaNs from zero division.
            dim (int, optional): Dimension to take median of. Defaults to 0.

        Returns:
            Tensor: Same shape as x, except dimension dim reduced.
        """
        # uncomment this assert for safety but might impact performance
        # assert (
        #     mask.sum(dim=dim).ne(0).all()
        # ), "mask should not be all False in any column, causes zero division"
        x_nan = x.float().masked_fill(~mask, float("nan"))
        if dim is None:
            x_median = x_nan.nanmedian()
        else:
            x_median, _ = x_nan.nanmedian(dim=dim, keepdim=keepdim)
        return x_median

    # Shrinking operator
    def shrink_lp_op(x: Tensor, beta: float, lp_norm: float) -> Tensor:
        if lp_norm == 1:
            return torch.sign(x) * torch.nn.functional.relu(torch.abs(x) - 1.0 / beta)
        else:
            return torch.sign(x) * torch.nn.functional.relu(
                torch.abs(x) - (1.0 / beta) * torch.pow(torch.abs(x), lp_norm - 1))