quantizers 1.0.0__py3-none-any.whl

quantizers/__init__.py

@@ -0,0 +1,4 @@
+from .binary import *
+from .fixed_point import *
+from .minifloat import *
+from .stateful_quantizers import BinaryQ, FixedQ, MinifloatQ, TernaryQ

quantizers/_version.py

@@ -0,0 +1,16 @@
+# file generated by setuptools_scm
+# don't change, don't track in version control
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple, Union
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+else:
+    VERSION_TUPLE = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+
+__version__ = version = '1.0.0'
+__version_tuple__ = version_tuple = (1, 0, 0)

quantizers/binary/__init__.py

@@ -0,0 +1,9 @@
+from .binary_ops_np import binary_quantize_np, ternary_quantize_np
+
+try:
+    from .binary_ops import binary_quantize, ternary_quantize
+except ImportError:
+    binary_quantize = binary_quantize_np
+    ternary_quantize = ternary_quantize_np
+
+__all__ = ['binary_quantize', 'ternary_quantize', 'binary_quantize_np', 'ternary_quantize_np']

quantizers/binary/binary_ops.py

@@ -0,0 +1,28 @@
+from keras import ops
+from keras.api.random import SeedGenerator
+
+
+@ops.custom_gradient
+def binary_quantize(x):
+    r = ops.where(x > 0, 1., -1.)
+
+    def grad(*args, upstream=None):
+        if upstream is None:
+            (upstream,) = args
+        dy = upstream
+        return dy * (1. - ops.tanh(x)**2)  # type: ignore
+
+    return ops.stop_gradient(r), grad
+
+
+@ops.custom_gradient
+def ternary_quantize(x):
+    r = ops.where(x > 0.5, 1., ops.where(x < -0.5, -1., 0.))
+
+    def grad(*args, upstream=None):
+        if upstream is None:
+            (upstream,) = args
+        dy = upstream
+        return dy * (1. - ops.tanh(x)**2)  # type: ignore
+
+    return ops.stop_gradient(r), grad

quantizers/binary/binary_ops_np.py

@@ -0,0 +1,9 @@
+import numpy as np
+
+
+def binary_quantize_np(x):
+    return np.where(x > 0, 1., -1.)
+
+
+def ternary_quantize_np(x):
+    return np.where(x > 0.5, 1., np.where(x < -0.5, -1., 0.))

quantizers/fixed_point/__init__.py

@@ -0,0 +1,8 @@
+from .fixed_point_ops_np import get_fixed_quantizer_np
+
+try:
+    from .fixed_point_ops import get_fixed_quantizer
+except ImportError:
+    get_fixed_quantizer = get_fixed_quantizer_np
+
+__all__ = ['get_fixed_quantizer', 'get_fixed_quantizer_np']

quantizers/fixed_point/fixed_point_ops.py

@@ -0,0 +1,279 @@
+from collections.abc import Callable
+from functools import singledispatch
+from typing import Any, TypeVar
+
+import keras
+from keras import ops
+from keras.api.random import SeedGenerator
+
+round_mode_registry: dict[str, Callable[[Any, bool | None], Any]] = {}
+round_mode_registry_scaled: dict[str, Callable[[Any, Any, bool | None, SeedGenerator | None], Any]] = {}
+saturation_mode_registry: dict[str, Callable[[Any, Any, Any, Any, bool | None], Any]] = {}
+from numpy.typing import ArrayLike
+
+T = TypeVar('T', bound=ArrayLike)
+
+
+@ops.custom_gradient
+def _clip(x, min_value, max_value):
+    r = ops.clip(x, min_value, max_value)
+
+    def grad(*args, upstream=None):
+        mask1 = x > max_value
+        mask2 = x < min_value
+        if upstream is None:
+            (upstream,) = args
+        dy = upstream
+        dx = ops.where(~(mask1 | mask2), dy, 0.)
+        dmax = ops.where(mask1, dy, 0.)
+        dmin = ops.where(mask2, dy, 0.)
+        return dx, dmin, dmax
+    return r, grad
+
+
+def rnd_mode(names: str | list[str] | tuple[str, ...]):
+    names = (names,) if isinstance(names, str) else names
+
+    def inner(func):
+        @keras.ops.custom_gradient
+        def inner_wrapper(x, f):
+            scale = 2.**f
+            sx = x * scale
+            sxq = func(sx)
+            xq = sxq / scale
+            delta = xq - x
+
+            def grad(*args, upstream=None):
+                if upstream is None:
+                    (upstream,) = args
+                dy = upstream
+                dx = dy
+                df = - ops.log(2.) * delta * dy  # type: ignore
+                return dx, df
+            return ops.stop_gradient(xq), grad
+
+        @keras.ops.custom_gradient
+        def inner_ste_wrapper(x):
+            r = func(x)
+
+            def grad(*args, upstream=None):
+                if upstream is None:
+                    (upstream,) = args
+                dy = upstream
+                return dy
+            return ops.stop_gradient(r), grad
+
+        def wrapper(x, f, training=None, seed_gen=None):  # type: ignore
+            return inner_wrapper(x, f)
+
+        def ste_wrapper(x, training=None):  # type: ignore
+            return inner_ste_wrapper(x)
+
+        for name in names:
+            assert name not in round_mode_registry_scaled, f"Rounding mode '{name}' already exists."
+            round_mode_registry_scaled[name] = wrapper
+        round_mode_registry_scaled[func.__name__.upper()] = wrapper
+
+        for name in names:
+            assert name not in round_mode_registry, f"Rounding mode '{name}' already exists."
+            round_mode_registry[name] = ste_wrapper
+        round_mode_registry[func.__name__.upper()] = ste_wrapper
+
+        return ste_wrapper
+    return inner
+
+
+@rnd_mode('TRN')
+def floor(x):
+    return ops.floor(x)
+
+
+@rnd_mode('RND')
+def round(x):
+    # Round to nearest, ties positive infinity.
+    return ops.floor(x + 0.5)
+
+
+@rnd_mode('RND_CONV')
+def round_conv(x):
+    # Round to nearest, ties to even.
+    return ops.round(x)
+
+
+@rnd_mode('TRN_ZERO')
+def floor_zero(x):
+    # Truncate towards zero.
+    sign = ops.sign(x)
+    return ops.floor(ops.abs(x)) * sign  # type: ignore
+
+
+@rnd_mode('RND_ZERO')
+def round_zero(x):
+    # Round to nearest, ties towards zero.
+    sign = ops.sign(x)
+    return -ops.floor(-ops.abs(x) + 0.5) * sign  # type:ignore
+
+
+@rnd_mode('RND_MIN_INF')
+def round_min_inf(x):
+    # Round to nearest, ties towards negative infinity.
+    return -ops.floor(-x + 0.5)  # type:ignore
+
+
+@rnd_mode('RND_INF')
+def round_inf(x):
+    # Round to nearest, ties away from zero.
+    sign = ops.sign(x)
+    return ops.floor(ops.abs(x) + 0.5) * sign  # type: ignore
+
+
+def sat_mode(name: str | list | tuple):
+    names = (name,) if isinstance(name, str) else name
+
+    def inner(func):
+        for name in names:
+            assert name not in saturation_mode_registry, f"Saturation mode '{name}' already exists."
+            saturation_mode_registry[name] = func
+        saturation_mode_registry[func.__name__.upper()] = func
+        return func
+    return inner
+
+
+@sat_mode('WRAP')
+def wrap(x, k, i, f, training=None):
+    if training:
+        return x
+
+    xs = x
+    bk = i + k
+    bias = k * 2**(bk - 1)
+    return ((xs + bias) % (2**bk) - bias)
+
+
+@sat_mode('SAT')
+def sat(x, k, i, f, training=None):
+    f_eps = 2**(-f)
+    __max = 2**i
+    _max = __max - f_eps
+    _min = -__max * k
+    r = _clip(x, _min, _max)
+    return r
+
+
+@sat_mode('SAT_SYM')
+def sat_sym(x, k, i, f, training=None):
+    f_eps = 2**(-f)
+    _max = 2**i - f_eps
+    _min = -_max * k
+    r = _clip(x, _min, _max)
+    return r
+
+
+@sat_mode('WRAP_SM')
+def wrap_sm_fn(x, k, i, f, training=None, quant_fn: Callable = lambda x: x):
+    # x=ops.round(x*2**f)
+    # High and low bounds are reflective. When overflows, can be less trash than WARP but still more trash than SAT.
+    eps = 2**-f
+    high = 2**i - eps
+    low = -(high + eps) * k
+    interval = 2 ** (i + k)
+    c1 = ((x) / interval) % 2 >= 1  # type: ignore
+    c1 = c1 & (ops.abs(x) > eps / 2)
+    c2 = ((x + eps / 2) / (interval / 2)) % 2 >= 1  # type: ignore
+    qx = quant_fn(x)
+    mapped = ((qx - high - eps) % interval) + low
+
+    mapped = ops.where(c2, -mapped - eps, mapped)  # type: ignore
+    mapped = ops.where(c1, -mapped - eps, mapped)  # type: ignore
+
+    return mapped
+
+
+def get_fixed_quantizer(round_mode: str = 'TRN', overflow_mode: str = 'WRAP') -> Callable[[T, Any, Any, Any, bool | None, SeedGenerator | None], T]:
+    """Get a stateless fixed-point quantizer given the round and overflow mode.
+    The quantizer is differentiable w.r.t. to the input and f, also i if using saturation overflow mode.
+
+    Args:
+        round_mode: round mode, one of
+    """
+    round_mode = round_mode.upper()
+    overflow_mode = overflow_mode.upper()
+    round_fn_scaled = round_mode_registry_scaled.get(round_mode, None)
+    if round_fn_scaled is None:
+        raise ValueError(f"Unknown rounding mode: {round_mode}")  # pragma: no cover
+    sat_fn = saturation_mode_registry.get(overflow_mode, None)
+    if sat_fn is None:
+        raise ValueError(f"Unknown saturation mode: {overflow_mode}")  # pragma: no cover
+
+    if overflow_mode == 'WRAP_SM':
+        assert round_mode in ('RND', 'RND_CONV'), "WRAP_SM only supports RND and RND_CONV rounding modes in this implementation."
+
+    def quantizer(x: T, k: T, i: T, f: T, training: bool | None = None, seed_gen: SeedGenerator | None = None) -> T:
+        """Stateless fixed-point quantizer.
+        Args:
+            x: input tensor
+            k: number of fractional bits
+            i: number of integer bits
+            f: number of fractional bits
+            training: training mode
+        """
+        i = ops.stop_gradient(ops.maximum(i, -f)) + (i - ops.stop_gradient(i))  # type: ignore
+
+        if overflow_mode == 'WRAP_SM':
+            def rnd_fn_wrapped(x):
+                return round_fn_scaled(x, f, training, seed_gen)
+            return wrap_sm_fn(x, k, i, f, training, rnd_fn_wrapped)  # type: ignore
+
+        # Workaround for gradient computation around 0.
+        # When have small rounded to boundary, grad on f presents despite the value will be clipped off anyway.
+        # Thus have saturation before rounding, except for wrap mode, which doesn't round during training.
+        if overflow_mode != 'WRAP':
+            x = sat_fn(x, k, i, f, training)
+        x = round_fn_scaled(x, f, training, seed_gen)
+        if overflow_mode == 'WRAP':
+            x = sat_fn(x, k, i, f, training)
+        return x
+    return quantizer
+
+
+# ======================= STOCHASTIC ROUNDING =======================
+def register_stochastic_rounding():
+    @keras.ops.custom_gradient
+    def stochastic_scaled(x, f, noise):
+        scale = 2.**f
+        sx = x * scale
+        sxq = ops.floor(sx + noise)
+        xq = sxq / scale
+        delta = xq - x
+
+        def grad(*args, upstream=None):
+            if upstream is None:
+                (upstream,) = args
+            dy = upstream
+            dx = dy
+            df = - ops.log(2.) * delta * dy  # type: ignore
+            return dx, df, None
+        return ops.stop_gradient(xq), grad
+
+    rnd_scaled = round_mode_registry_scaled['RND']
+    rnd_conv_scaled = round_mode_registry_scaled['RND_CONV']
+
+    def outer_rnd(x, f, training=None, seed_gen=None):
+        if training:
+            noise = keras.random.uniform(ops.shape(x), seed=seed_gen)
+            return stochastic_scaled(x, f, noise)
+        else:
+            return rnd_scaled(x, f, training, seed_gen)
+
+    def outer_rnd_conv(x, f, training=None, seed_gen=None):
+        if training:
+            noise = keras.random.uniform(ops.shape(x), seed=seed_gen)
+            return stochastic_scaled(x, f, noise)
+        else:
+            return rnd_conv_scaled(x, f, training, seed_gen)
+
+    round_mode_registry_scaled['S_RND'] = outer_rnd
+    round_mode_registry_scaled['S_RND_CONV'] = outer_rnd_conv
+
+
+register_stochastic_rounding()

quantizers/fixed_point/fixed_point_ops_np.py

@@ -0,0 +1,163 @@
+from collections.abc import Callable
+from typing import Any, TypeVar
+
+import numpy as np
+from numpy.typing import ArrayLike
+
+round_mode_registry_np: dict[str, Callable[[Any, bool | None], Any]] = {}
+round_mode_registry_scaled_np: dict[str, Callable[[Any, Any, bool | None], Any]] = {}
+saturation_mode_registry_np: dict[str, Callable[[Any, Any, Any, Any, bool | None], Any]] = {}
+
+T = TypeVar('T', bound=np.ndarray)
+
+
+def rnd_mode_np(names: str | list[str] | tuple[str, ...]):
+    names = (names,) if isinstance(names, str) else names
+
+    def inner(func):
+        def wrapper(x, f, training=None, seed_gen=None):
+            scale = 2.**f
+            sx = x * scale
+            sxq = func(sx)
+            xq = sxq / scale
+            return xq
+
+        for name in names:
+            assert name not in round_mode_registry_scaled_np, f"Rounding mode '{name}' already exists."
+            round_mode_registry_scaled_np[name] = wrapper
+        round_mode_registry_scaled_np[func.__name__.upper()] = wrapper
+
+        for name in names:
+            assert name not in round_mode_registry_np, f"Rounding mode '{name}' already exists."
+            round_mode_registry_np[name] = func
+        round_mode_registry_np[func.__name__.upper()] = func
+
+        return func
+    return inner
+
+
+@rnd_mode_np('TRN')
+def floor(x):
+    return np.floor(x)
+
+
+@rnd_mode_np('RND')
+def round(x):
+    return np.floor(x + 0.5)
+
+
+@rnd_mode_np('RND_CONV')
+def round_conv(x):
+    return np.round(x)
+
+
+@rnd_mode_np('TRN_ZERO')
+def floor_zero(x):
+    sign = np.sign(x)
+    return np.floor(np.abs(x)) * sign
+
+
+@rnd_mode_np('RND_ZERO')
+def round_zero(x):
+    sign = np.sign(x)
+    return -np.floor(-np.abs(x) + 0.5) * sign
+
+
+@rnd_mode_np('RND_MIN_INF')
+def round_min_inf(x):
+    return -np.floor(-x + 0.5)
+
+
+@rnd_mode_np('RND_INF')
+def round_inf(x):
+    sign = np.sign(x)
+    return np.floor(np.abs(x) + 0.5) * sign
+
+
+def sat_mode_np(name: str | list | tuple):
+    names = (name,) if isinstance(name, str) else name
+
+    def inner(func):
+        for name in names:
+            assert name not in saturation_mode_registry_np, f"Saturation mode '{name}' already exists."
+            saturation_mode_registry_np[name] = func
+        saturation_mode_registry_np[func.__name__.upper()] = func
+        return func
+    return inner
+
+
+@sat_mode_np('WRAP')
+def wrap(x, k, i, f, training=None):
+    xs = x
+    bk = i + k
+    bias = k * 2**(bk - 1)
+    return ((xs + bias) % (2**bk) - bias)
+
+
+@sat_mode_np('SAT')
+def sat(x, k, i, f, training=None):
+    f_eps = 2**(-f)
+    __max = 2**i
+    _max = __max - f_eps
+    _min = -__max * k
+    return np.clip(x, _min, _max)
+
+
+@sat_mode_np('SAT_SYM')
+def sat_sym(x, k, i, f, training=None):
+    f_eps = 2**(-f)
+    _max = 2**i - f_eps
+    _min = -_max * k
+    return np.clip(x, _min, _max)
+
+
+@sat_mode_np('WRAP_SM')
+def wrap_sm_fn(x, k, i, f, training=None, quant_fn: Callable = lambda x: x):
+    eps = 2**-f
+    high = 2**i - eps
+    low = -(high + eps) * k
+    interval = 2 ** (i + k)
+    c1 = ((x) / interval) % 2 >= 1
+    c1 = c1 & (np.abs(x) > eps / 2)
+    c2 = ((x + eps / 2) / (interval / 2)) % 2 >= 1
+    qx = quant_fn(x)
+    mapped = ((qx - high - eps) % interval) + low
+
+    mapped = np.where(c2, -mapped - eps, mapped)
+    mapped = np.where(c1, -mapped - eps, mapped)
+
+    return mapped
+
+
+def get_fixed_quantizer_np(round_mode: str = 'TRN', overflow_mode: str = 'WRAP'):
+    round_mode = round_mode.upper()
+    overflow_mode = overflow_mode.upper()
+    round_mode = round_mode[2:] if round_mode.startswith('S_') else round_mode
+
+    round_fn_scaled = round_mode_registry_scaled_np.get(round_mode, None)
+    if round_fn_scaled is None:
+        raise ValueError(f"Unknown rounding mode: {round_mode}")  # pragma: no cover
+    sat_fn = saturation_mode_registry_np.get(overflow_mode, None)
+    if sat_fn is None:
+        raise ValueError(f"Unknown saturation mode: {overflow_mode}")  # pragma: no cover
+
+    if overflow_mode == 'WRAP_SM':
+        assert round_mode in ('RND', 'RND_CONV'), "WRAP_SM only supports RND and RND_CONV rounding modes in this implementation."
+
+    def quantizer(x: T, k: T, i: T, f: T, training: bool | None = False, seed_gen: None = None) -> np.ndarray:
+        assert not training, "Training mode not supported in numpy implementation."
+
+        _i = np.maximum(i, -f)
+
+        if overflow_mode == 'WRAP_SM':
+            def rnd_fn_wrapped(x):
+                return round_fn_scaled(x, f, training)
+            return wrap_sm_fn(x, k, _i, f, training, rnd_fn_wrapped)
+
+        if overflow_mode != 'WRAP':
+            x = sat_fn(x, k, _i, f, training)
+        x = round_fn_scaled(x, f, training)
+        if overflow_mode == 'WRAP':
+            x = sat_fn(x, k, _i, f, training)
+        return x
+    return quantizer

quantizers/minifloat/__init__.py

@@ -0,0 +1,9 @@
+from .float_point_ops_np import float_decompose_np, float_quantize_np
+
+try:
+    from .float_point_ops import float_decompose, float_quantize
+except ImportError:
+    float_quantize = float_quantize_np
+    float_decompose = float_decompose_np
+
+__all__ = ['float_quantize', 'float_decompose', 'float_quantize_np', 'float_decompose_np']

quantizers/minifloat/float_point_ops.py

@@ -0,0 +1,56 @@
+import numpy as np
+from keras import ops
+
+
+@ops.custom_gradient
+def float_quantize(x, m, e, e0=0.):
+    '''Quantize an array to floatlet (m mantissa bits, excl. sign bit, e exponent bits) format. Tentative gradient impl.'''
+    m = m + 1
+    eps = 1e-7
+
+    e_req = ops.floor(ops.log2(ops.abs(x) + eps))  # type: ignore
+    _e_high = ops.maximum(2.**(e - 1), 1.)
+    e_low, e_high = -_e_high + e0 + 1, _e_high + e0 - 1  # type: ignore
+    e_act = ops.clip(e_req, e_low, e_high)
+    scale = 2.**(e_act - m + 1)
+    sig = x / scale
+    qsig = ops.round(sig)
+    _sig_high = 2.**m
+    _sig_high = ops.where(e_high != e_act, _sig_high, _sig_high - 1)
+    clip_sig = ops.clip(qsig, -_sig_high, _sig_high)  # type: ignore
+    qx = clip_sig * scale
+
+    def grad(*args, upstream=None):
+        if upstream is None:
+            (upstream,) = args
+        dy = upstream
+        mask_e = e_req != e_act
+        mask_m = e_req <= e_high  # type: ignore
+        dm = ops.where(mask_m, scale * (sig - qsig) * dy * ops.log(2.), 0.)
+        d_exp = ops.where(mask_e, (x - qx) * dy * ops.log(2.), 0.)
+        de = d_exp * ops.log(2.) * _e_high  # type: ignore
+        de0 = ops.where(e_req > e_high, d_exp, -d_exp)  # type: ignore
+        return dy, dm, de, de0
+
+    return ops.stop_gradient(qx), grad
+
+
+def float_decompose(x, m, e, e0=0):
+    '''Quantize an array to floatlet (m mantissa bits, excl. sign bit, e exponent bits) format. Tentative gradient impl.'''
+    m = m + 1
+    eps = 1e-7
+
+    e_req = ops.floor(ops.log2(ops.abs(x) + eps))  # type: ignore
+    _e_high = 2.**(e - 1)
+    e_low, e_high = -_e_high + e0, _e_high + e0 - 1
+    e_act = ops.clip(e_req, e_low + 1, e_high)
+    scale = 2.**(e_act - m + 1)
+    sig = x / scale
+    qsig = ops.round(sig)
+    cond = (e_act != e_high) & (ops.abs(qsig) == 2**m)
+    e_act = ops.where(cond, e_act + 1, e_act)  # type: ignore
+    qsig = ops.where(cond, qsig / 2., qsig)  # type: ignore
+    r_mantissa_bound = 2 - 2.**(-m + 1)
+    r_mantissa = ops.clip(qsig * 2.**(-m + 1), -r_mantissa_bound, r_mantissa_bound)
+    r_exponent = e_act
+    return r_mantissa, r_exponent

quantizers/minifloat/float_point_ops_np.py

@@ -0,0 +1,42 @@
+import numpy as np
+
+
+def float_quantize_np(x, m, e, e0=0.):
+    '''Quantize an array to floatlet (m mantissa bits, excl. sign bit, e exponent bits) format. Tentative gradient impl.'''
+    m = m + 1
+    eps = 1e-7
+
+    e_req = np.floor(np.log2(np.abs(x) + eps))  # type: ignore
+    _e_high = np.maximum(2.**(e - 1), 1.)
+    e_low, e_high = -_e_high + e0 + 1, _e_high + e0 - 1  # type: ignore
+    e_act = np.clip(e_req, e_low, e_high)
+    scale = 2.**(e_act - m + 1)
+    sig = x / scale
+    qsig = np.round(sig)
+    _sig_high = 2.**m
+    _sig_high = np.where(e_high != e_act, _sig_high, _sig_high - 1)
+    clip_sig = np.clip(qsig, -_sig_high, _sig_high)  # type: ignore
+    qx = clip_sig * scale
+
+    return qx
+
+
+def float_decompose_np(x, m, e, e0=0):
+    '''Quantize an array to floatlet (m mantissa bits, excl. sign bit, e exponent bits) format. Tentative gradient impl.'''
+    m = m + 1
+    eps = 1e-7
+
+    e_req = np.floor(np.log2(np.abs(x) + eps))  # type: ignore
+    _e_high = 2.**(e - 1)
+    e_low, e_high = -_e_high + e0, _e_high + e0 - 1
+    e_act = np.clip(e_req, e_low + 1, e_high)
+    scale = 2.**(e_act - m + 1)
+    sig = x / scale
+    qsig = np.round(sig)
+    cond = (e_act != e_high) & (np.abs(qsig) == 2**m)
+    e_act = np.where(cond, e_act + 1, e_act)  # type: ignore
+    qsig = np.where(cond, qsig / 2., qsig)  # type: ignore
+    r_mantissa_bound = 2 - 2.**(-m + 1)
+    r_mantissa = np.clip(qsig * 2.**(-m + 1), -r_mantissa_bound, r_mantissa_bound)
+    r_exponent = e_act
+    return r_mantissa, r_exponent

quantizers/stateful_quantizers.py

@@ -0,0 +1,142 @@
+from typing import Any, TypeVar, overload
+
+import numpy as np
+
+from .binary import *
+from .fixed_point import *
+from .minifloat import *
+
+T = TypeVar('T', np.ndarray, float, int)
+
+
+def is_tensor(x):
+    base_module = x.__class__.__module__.split('.', 1)[0]
+    return base_module in ('tensorflow', 'torch', 'jaxlib')
+
+
+class BinaryQ:
+    """Stateful binary quantizer."""
+
+    @overload
+    def __call__(self, x: T) -> np.ndarray:
+        ...  # pragma: no cover
+
+    @overload
+    def __call__(self, x: Any) -> Any:
+        ...  # pragma: no cover
+
+    def __call__(self, x):
+        use_numpy = not is_tensor(x)
+        if use_numpy:
+            return binary_quantize_np(x)
+        else:
+            return binary_quantize(x)
+
+    def __repr__(self):
+        return 'Binary()'
+
+
+class TernaryQ:
+    """Stateful ternary quantizer."""
+
+    @overload
+    def __call__(self, x: T) -> np.ndarray:
+        ...  # pragma: no cover
+
+    @overload
+    def __call__(self, x: Any) -> Any:
+        ...  # pragma: no cover
+
+    def __call__(self, x):
+        use_numpy = not is_tensor(x)
+        if use_numpy:
+            return ternary_quantize_np(x)
+        else:
+            return ternary_quantize(x)
+
+    def __repr__(self):
+        return 'Ternary()'
+
+
+class FixedQ:
+    """Stateful fixed-point quantizer.
+
+    Parameters
+    ----------
+    width : int
+        The total bit-width of the quantized number.
+    integer : int
+        The number of bits used for the integer part. INCLUDES the sign bit.
+    signed : bool|int
+        Whether the quantized number is signed.
+    rounding : str, optional
+        The rounding method to use, by default 'TRN'. Available options are 'TRN', 'TRN_ZERO', 'RND', 'S_RND', 'S_RND_CONV', 'RND_CONV', 'RND_ZERO', 'RND_INF', 'RND_MIN_INF'.
+    overflow : str, optional
+        The overflow handling method to use, by default 'WRAP'. Available options are 'WRAP', 'SAT', 'SAT_SYM', 'WRAP_SM'.
+    """
+
+    def __init__(self, width: int, integer: int, signed: bool | int = True, rounding: str = 'TRN', overflow: str = 'WRAP'):
+        self.width = width
+        self.integer = integer
+        self.signed = signed
+        self.rounding = rounding
+        self.overflow = overflow
+
+    @overload
+    def __call__(self, x: T) -> np.ndarray:
+        ...  # pragma: no cover
+
+    @overload
+    def __call__(self, x: Any) -> Any:
+        ...  # pragma: no cover
+
+    def __call__(self, x):
+        use_numpy = not is_tensor(x)
+        k, i, f = self.signed, self.integer - self.signed, self.width - self.integer
+        k, i, f = float(k), float(i), float(f)
+        if use_numpy:
+            q = get_fixed_quantizer_np(self.rounding, self.overflow)
+        else:
+            q = get_fixed_quantizer(self.rounding, self.overflow)
+        return q(x, k, i, f, False, None)  # type: ignore
+
+    def __repr__(self):
+        return f'Fixed(w={self.width}, i={self.integer}, s={self.signed}, r={self.rounding}, o={self.overflow})'
+
+
+class MinifloatQ:
+    """Stateful minifloat quantizer.
+
+    Parameters
+    ----------
+    significant : int
+        The number of significant bits.
+    exponent : int
+        The number of exponent bits.
+    exponent_bias : int
+        The bias applied to the exponent.
+    """
+
+    def __init__(self, significant: int, exponent: int, exponent_bias: int):
+        self.significant = significant
+        self.exponent = exponent
+        self.exponent_bias = exponent_bias
+
+    @overload
+    def __call__(self, x: T) -> np.ndarray:
+        ...  # pragma: no cover
+
+    @overload
+    def __call__(self, x: Any) -> Any:
+        ...  # pragma: no cover
+
+    def __call__(self, x):
+        s, e, e0 = float(self.significant), float(self.exponent), float(self.exponent_bias)
+        use_numpy = not is_tensor(x)
+        if use_numpy:
+            return float_quantize_np(x, s, e, e0)
+        else:
+            return float_quantize(x, s, e, e0)
+
+    def __repr__(self):
+        return f'Minifloat(s={self.significant}, e={self.exponent}, b={self.exponent_bias})'

quantizers-1.0.0.dist-info/LICENSE

@@ -0,0 +1,165 @@
+                   GNU LESSER GENERAL PUBLIC LICENSE
+                       Version 3, 29 June 2007
+
+ Copyright (C) 2007 Free Software Foundation, Inc. <https://fsf.org/>
+ Everyone is permitted to copy and distribute verbatim copies
+ of this license document, but changing it is not allowed.
+
+
+  This version of the GNU Lesser General Public License incorporates
+the terms and conditions of version 3 of the GNU General Public
+License, supplemented by the additional permissions listed below.
+
+  0. Additional Definitions.
+
+  As used herein, "this License" refers to version 3 of the GNU Lesser
+General Public License, and the "GNU GPL" refers to version 3 of the GNU
+General Public License.
+
+  "The Library" refers to a covered work governed by this License,
+other than an Application or a Combined Work as defined below.
+
+  An "Application" is any work that makes use of an interface provided
+by the Library, but which is not otherwise based on the Library.
+Defining a subclass of a class defined by the Library is deemed a mode
+of using an interface provided by the Library.
+
+  A "Combined Work" is a work produced by combining or linking an
+Application with the Library.  The particular version of the Library
+with which the Combined Work was made is also called the "Linked
+Version".
+
+  The "Minimal Corresponding Source" for a Combined Work means the
+Corresponding Source for the Combined Work, excluding any source code
+for portions of the Combined Work that, considered in isolation, are
+based on the Application, and not on the Linked Version.
+
+  The "Corresponding Application Code" for a Combined Work means the
+object code and/or source code for the Application, including any data
+and utility programs needed for reproducing the Combined Work from the
+Application, but excluding the System Libraries of the Combined Work.
+
+  1. Exception to Section 3 of the GNU GPL.
+
+  You may convey a covered work under sections 3 and 4 of this License
+without being bound by section 3 of the GNU GPL.
+
+  2. Conveying Modified Versions.
+
+  If you modify a copy of the Library, and, in your modifications, a
+facility refers to a function or data to be supplied by an Application
+that uses the facility (other than as an argument passed when the
+facility is invoked), then you may convey a copy of the modified
+version:
+
+   a) under this License, provided that you make a good faith effort to
+   ensure that, in the event an Application does not supply the
+   function or data, the facility still operates, and performs
+   whatever part of its purpose remains meaningful, or
+
+   b) under the GNU GPL, with none of the additional permissions of
+   this License applicable to that copy.
+
+  3. Object Code Incorporating Material from Library Header Files.
+
+  The object code form of an Application may incorporate material from
+a header file that is part of the Library.  You may convey such object
+code under terms of your choice, provided that, if the incorporated
+material is not limited to numerical parameters, data structure
+layouts and accessors, or small macros, inline functions and templates
+(ten or fewer lines in length), you do both of the following:
+
+   a) Give prominent notice with each copy of the object code that the
+   Library is used in it and that the Library and its use are
+   covered by this License.
+
+   b) Accompany the object code with a copy of the GNU GPL and this license
+   document.
+
+  4. Combined Works.
+
+  You may convey a Combined Work under terms of your choice that,
+taken together, effectively do not restrict modification of the
+portions of the Library contained in the Combined Work and reverse
+engineering for debugging such modifications, if you also do each of
+the following:
+
+   a) Give prominent notice with each copy of the Combined Work that
+   the Library is used in it and that the Library and its use are
+   covered by this License.
+
+   b) Accompany the Combined Work with a copy of the GNU GPL and this license
+   document.
+
+   c) For a Combined Work that displays copyright notices during
+   execution, include the copyright notice for the Library among
+   these notices, as well as a reference directing the user to the
+   copies of the GNU GPL and this license document.
+
+   d) Do one of the following:
+
+       0) Convey the Minimal Corresponding Source under the terms of this
+       License, and the Corresponding Application Code in a form
+       suitable for, and under terms that permit, the user to
+       recombine or relink the Application with a modified version of
+       the Linked Version to produce a modified Combined Work, in the
+       manner specified by section 6 of the GNU GPL for conveying
+       Corresponding Source.
+
+       1) Use a suitable shared library mechanism for linking with the
+       Library.  A suitable mechanism is one that (a) uses at run time
+       a copy of the Library already present on the user's computer
+       system, and (b) will operate properly with a modified version
+       of the Library that is interface-compatible with the Linked
+       Version.
+
+   e) Provide Installation Information, but only if you would otherwise
+   be required to provide such information under section 6 of the
+   GNU GPL, and only to the extent that such information is
+   necessary to install and execute a modified version of the
+   Combined Work produced by recombining or relinking the
+   Application with a modified version of the Linked Version. (If
+   you use option 4d0, the Installation Information must accompany
+   the Minimal Corresponding Source and Corresponding Application
+   Code. If you use option 4d1, you must provide the Installation
+   Information in the manner specified by section 6 of the GNU GPL
+   for conveying Corresponding Source.)
+
+  5. Combined Libraries.
+
+  You may place library facilities that are a work based on the
+Library side by side in a single library together with other library
+facilities that are not Applications and are not covered by this
+License, and convey such a combined library under terms of your
+choice, if you do both of the following:
+
+   a) Accompany the combined library with a copy of the same work based
+   on the Library, uncombined with any other library facilities,
+   conveyed under the terms of this License.
+
+   b) Give prominent notice with the combined library that part of it
+   is a work based on the Library, and explaining where to find the
+   accompanying uncombined form of the same work.
+
+  6. Revised Versions of the GNU Lesser General Public License.
+
+  The Free Software Foundation may publish revised and/or new versions
+of the GNU Lesser General Public License from time to time. Such new
+versions will be similar in spirit to the present version, but may
+differ in detail to address new problems or concerns.
+
+  Each version is given a distinguishing version number. If the
+Library as you received it specifies that a certain numbered version
+of the GNU Lesser General Public License "or any later version"
+applies to it, you have the option of following the terms and
+conditions either of that published version or of any later version
+published by the Free Software Foundation. If the Library as you
+received it does not specify a version number of the GNU Lesser
+General Public License, you may choose any version of the GNU Lesser
+General Public License ever published by the Free Software Foundation.
+
+  If the Library as you received it specifies that a proxy can decide
+whether future versions of the GNU Lesser General Public License shall
+apply, that proxy's public statement of acceptance of any version is
+permanent authorization for you to choose that version for the
+Library.

quantizers-1.0.0.dist-info/METADATA

@@ -0,0 +1,131 @@
+Metadata-Version: 2.1
+Name: quantizers
+Version: 1.0.0
+Author-email: Chang Sun <chsun@cern.ch>
+License: GNU Lesser General Public License v3 (LGPLv3)
+Project-URL: repository, https://github.com/calad0i/quantizers
+Keywords: quantization
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: OS Independent
+Classifier: License :: OSI Approved :: GNU Lesser General Public License v3 (LGPLv3)
+Classifier: Development Status :: 5 - Production/Stable
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Provides-Extra: dependencies
+Requires-Dist: numpy ; extra == 'dependencies'
+Provides-Extra: grad
+Requires-Dist: keras >=3 ; extra == 'grad'
+Provides-Extra: test
+Requires-Dist: pytest ; extra == 'test'
+Requires-Dist: pytest-cov ; extra == 'test'
+Requires-Dist: pytest-xdist ; extra == 'test'
+Requires-Dist: keras >=3 ; extra == 'test'
+
+# Quantizers
+[![PyPI version](https://badge.fury.io/py/quantizers.svg)](https://badge.fury.io/py/quantizers)
+[![License](https://img.shields.io/badge/License-LGPL-blue)](LICENSE)
+
+Hardware-oriented numerical quantizers for deep learning models, implemented in Keras v3 and NumPy. Provides bit-accurate precision matching with Vivado/Vitis HLS implementations.
+
+## Features
+
+- Bit-accurate to the HLS implementation up to 32/64-bit floating point precision
+- Support for fixed-point and minifloat number formats
+- Differentiable Keras v3 implementations with gradients on inputs
+  - With surrogate gradients for bit-width optimization as described in *[Gradient-based Automatic Mixed Precision Quantization for Neural Networks On-Chip](https://arxiv.org/abs/2405.00645)*
+- Supports stochastic rounding for training
+
+## Supported Quantizers
+
+### Fixed-Point Quantizer
+
+Parameters:
+- `k` (keep_negative): Enable negative numbers
+- `i` (integer_bits): Number of bits before decimal point (excludes sign bit)
+- `f` (fractional_bits): Number of bits after decimal point
+- For C++: `W = k + i + f`, `I = k + i`, `S = k`
+
+Supported modes:
+- Rounding: `TRN`, `RND`, `RND_CONV`, `TRN_ZERO`, `RND_ZERO`, `RND_MIN_INF`, `RND_INF`
+  - `S_RND` and `S_RND_CONV` for stochastic rounding; Not available in NumPy implementation as it is for training only
+- Overflow: `WRAP`, `SAT`, `SAT_SYM`, `WRAP_SM`
+
+Limitations:
+- `WRAP_SM` only works with `RND` or `RND_CONV` rounding
+- `WRAP*` modes don't provide surrogate gradients for integer bits
+- Saturation bit forced to zero for `WRAP` and `WRAP_SM`
+
+### Minifloat Quantizer
+
+Parameters:
+- `m` (mantissa_bits): Mantissa width
+- `e` (exponent_bits): Exponent width
+- `e0` (exponent_zero): Exponent bias (default: 0)
+- Range: `[-2^(e-1) + e0, 2^(e-1) - 1 + e0]`
+
+Features:
+- Supports subnormal numbers
+- Uses `RND_CONV` rounding and `SAT` overflow
+- HLS-synthesizable implementation in `test/cpp_source/ap_types/ap_float.h`
+
+### Simplified Quantizers
+
+- **Binary**: Maps to {-1,1} with 0 to -1. (preliminary implementation)
+- **Ternary**: Shorthand for fixed-point `fixed<2, 1, RND_CONV, SAT_SYM>`
+
+
+## Installation
+
+**requires python>=3.10**
+
+```bash
+pip install quantizers
+```
+`keras>=3.0` and at least one compatible backend (`pytorch`, `jax`, or `tensorflow`) is required for training.
+
+## Usage
+
+### Stateless Quantizers
+```python
+from quantizers import (
+  float_quantize(_np), # add _np for NumPy implementation
+  get_fixed_quantizer(_np),
+  binary_quantize(_np),
+  ternary_quantize(_np),
+)
+
+# Fixed-point quantizer
+fixed_quantizer = get_fixed_quantizer(round_mode, overflow_mode)
+fixedp_qtensor = fixed_quantizer(
+    x,
+    integer_bits,
+    fractional_bits,
+    keep_negative,
+    training, # For stochastic rounding, and WRAP does not happen during training
+    seed, # For stochastic rounding only
+)
+
+# Minifloat quantizer
+floatp_qtensor = float_quantize(x, mantissa_bits, exponent_bits, exponent_zero)
+
+# Simplified quantizers
+binary_qtensor = binary_quantize(x)
+ternary_qtensor = ternary_quantize(x)
+```
+
+### Stateful Quantizers
+```python
+# Can be used for, but not intended for training
+fixed_q = FixedQ(
+    width,
+    integer_bits, # including the sign bit)
+    keep_negative,
+    fixed_round_mode, # No stochastic rounding
+    fixed_overflow_mode
+)
+quantized = fixed_q(x)
+
+mfloat_q = MinifloatQ(mantissa_bits, exponent_bits, exponent_zero)
+quantized = mfloat_q(x)
+```

quantizers-1.0.0.dist-info/RECORD

@@ -0,0 +1,17 @@
+quantizers/__init__.py,sha256=o6yFBm_3iSB2Lz2oRTQHYA5o7s49i5La7Ib66P0bxXU,145
+quantizers/_version.py,sha256=DGJ4pj32xs3_DRJhSzQwCiRNnAQrMgo09USYpyMZsKc,411
+quantizers/stateful_quantizers.py,sha256=B7VuDJnpMR7Pp9bPGoq_FGyx8ykLx9B-_6dDKLnRdhA,4007
+quantizers/binary/__init__.py,sha256=tguWtPkpGhAVzAo28fdh1SGc7Bk3Ww33FBU0K829MHo,335
+quantizers/binary/binary_ops.py,sha256=xYFHN-3aENv3swwn7YQIsW9xsLbXXbJXZuCXPYEV07o,687
+quantizers/binary/binary_ops_np.py,sha256=GMjT9BUmkN0ykkLvQGvyh8k6Lj5x-IQRmy92qm0phvo,176
+quantizers/fixed_point/__init__.py,sha256=X7zxaHzTG7uFfW0warVxef2WMvVYtiqNH-Lo3HlcPPo,244
+quantizers/fixed_point/fixed_point_ops.py,sha256=QDsyQYCCuk-FK9nk51Sa1GOuLEidQ8LFwKanLK4a9Xo,8900
+quantizers/fixed_point/fixed_point_ops_np.py,sha256=3oHsEB_R_20u6VrvoPH3Ahx6xf234Xs9hX2S4jQha9A,4748
+quantizers/minifloat/__init__.py,sha256=6AZVKSRESofllv0joQR6KAs4w9Eay8f6UAkE89ieW5s,333
+quantizers/minifloat/float_point_ops.py,sha256=xOrqoHtPCBb2dJnbiMX3mMyf2JsPUKXGs96JGJeIezg,2105
+quantizers/minifloat/float_point_ops_np.py,sha256=e2t_cOjuzwd_JGKEmHZj-POY-SLiXOtPKbswzJfkndc,1522
+quantizers-1.0.0.dist-info/LICENSE,sha256=46mU2C5kSwOnkqkw9XQAJlhBL2JAf1_uCD8lVcXyMRg,7652
+quantizers-1.0.0.dist-info/METADATA,sha256=gg-4sGI-3RtD53Tr1APALE774MXO1BtKJSK07KxEZyA,4306
+quantizers-1.0.0.dist-info/WHEEL,sha256=R06PA3UVYHThwHvxuRWMqaGcr-PuniXahwjmQRFMEkY,91
+quantizers-1.0.0.dist-info/top_level.txt,sha256=K9nZlz267J8tCfY-mMGqqb8umqII1BxWSIXgC4J-vn0,11
+quantizers-1.0.0.dist-info/RECORD,,

quantizers-1.0.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (75.5.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

quantizers-1.0.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+quantizers