quantizers 1.1.0__tar.gz → 1.2.0__tar.gz

{quantizers-1.1.0/src/quantizers.egg-info → quantizers-1.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.4
 Name: quantizers
-Version: 1.1.0
+Version: 1.2.0
 Author-email: Chang Sun <chsun@cern.ch>
 License: GNU Lesser General Public License v3 (LGPLv3)
 Project-URL: repository, https://github.com/calad0i/quantizers
@@ -24,6 +24,7 @@ Requires-Dist: pytest-env; extra == "test"
 Requires-Dist: keras>=3.0; extra == "test"
 Requires-Dist: cppyy; extra == "test"
 Requires-Dist: jax; extra == "test"
+Dynamic: license-file
 
 # Quantizers
 [![PyPI version](https://badge.fury.io/py/quantizers.svg)](https://badge.fury.io/py/quantizers)

{quantizers-1.1.0 → quantizers-1.2.0}/pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "quantizers"
-version = "1.1.0"
+version = "1.2.0"
 authors = [{ name = "Chang Sun", email = "chsun@cern.ch" }]
 description = ""
 readme = "README.md"
@@ -91,7 +91,7 @@ docstring-code-line-length = 130
 
 [tool.ruff.lint]
 select = ["E", "W", "F", "I", "F401"]
-ignore = ["E501", "F403", "F405"]
+ignore = ["E501", "F403", "F405", "E741"]
 explicit-preview-rules = true
 fixable = ["ALL"]
 unfixable = []

{quantizers-1.1.0 → quantizers-1.2.0}/src/quantizers/_version.py

@@ -1,8 +1,13 @@
-# file generated by setuptools_scm
+# file generated by setuptools-scm
 # don't change, don't track in version control
+
+__all__ = ["__version__", "__version_tuple__", "version", "version_tuple"]
+
 TYPE_CHECKING = False
 if TYPE_CHECKING:
-    from typing import Tuple, Union
+    from typing import Tuple
+    from typing import Union
+
     VERSION_TUPLE = Tuple[Union[int, str], ...]
 else:
     VERSION_TUPLE = object
@@ -12,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '1.1.0'
-__version_tuple__ = version_tuple = (1, 1, 0)
+__version__ = version = '1.2.0'
+__version_tuple__ = version_tuple = (1, 2, 0)

quantizers-1.2.0/src/quantizers/fixed_point/_fixed_point_ops.py

@@ -0,0 +1,215 @@
+from collections.abc import Callable
+from functools import wraps
+from typing import Any, TypeVar
+
+from keras import ops
+from keras.api.random import SeedGenerator
+from numpy.typing import ArrayLike
+
+round_mode_registry: dict[str, Callable[[Any], Any]] = {}
+saturation_mode_registry: dict[str, Callable[[Any, Any, Any, Any], Any]] = {}
+
+T = TypeVar('T', bound=ArrayLike)
+
+
+def _clip(x, min_value, max_value):
+    mask_overflow = x > max_value
+    mask_underflow = x < min_value
+    return ops.where(mask_overflow, max_value, ops.where(mask_underflow, min_value, x))
+
+
+def rnd_mode(name: str):
+    def inner(func):
+        assert name not in round_mode_registry, f"Round mode '{name}' already exists."
+
+        @wraps(func)
+        def wrapper(x):
+            xq = func(x)
+            return ops.stop_gradient(xq) + (x - ops.stop_gradient(x))
+
+        round_mode_registry[name] = wrapper
+        return 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):
+    xs = x
+    bk = i + k
+    bias = k * 2.0 ** (bk - 1)
+    return (xs + bias) % (2.0**bk) - bias
+
+
+@sat_mode('SAT')
+def sat(x, k, i, f):
+    f_eps = 2.0 ** (-f)
+    __max = 2.0**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):
+    f_eps = 2.0 ** (-f)
+    _max = 2.0**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.0**-f
+    high = 2.0**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
+
+
+class FixedPointQuantizer:
+    def round(self, x, f: Any = 1.0, stochastic: bool | None = None, seed_gen: SeedGenerator | None = None):
+        scale = 2.0**f
+        x = x * scale
+        xq = self.round_fn(x)
+        # if stochastic:
+        #     noise = keras.random.uniform(ops.shape(x), -0.49, 0.49, seed=seed_gen)
+        #     x = x + noise
+        #     xq2 = self.round_fn(x)
+        #     xq = xq + ops.stop_gradient(xq2 - xq)
+        xq = xq / scale
+        return xq
+
+    def saturate(self, x, k, i, f):
+        return self.sat_fn(x, k, i, f)
+
+    def __init__(self, round_mode: str = 'TRN', overflow_mode: str = 'WRAP'):
+        round_mode = round_mode.upper()
+        overflow_mode = overflow_mode.upper()
+        self.stochastic = False
+
+        if round_mode.startswith('S_'):
+            round_mode = round_mode[2:]
+            self.stochastic = True
+
+        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.'
+
+        self.round_mode = round_mode
+        self.overflow_mode = overflow_mode
+        round_fn = round_mode_registry[round_mode]
+        sat_fn = saturation_mode_registry[overflow_mode]
+        self.sat_fn = sat_fn
+        self.round_fn = round_fn
+
+    def forward(self, x, k, i, f, training=False, seed_gen=None):
+        # Workaround for gradient computation around 0.
+        # When have values outside boundary 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 self.overflow_mode != 'WRAP':
+            x = self.saturate(x, k, i, f)
+        x = self.round(x, f, self.stochastic and training, seed_gen)
+        if self.overflow_mode == 'WRAP' and not training:
+            x = self.saturate(x, k, i, f)
+        return x
+
+    def forward_wrap_sm(self, x, k, i, f, training=False, seed_gen: SeedGenerator | None = None):
+        def quant_fn(x):
+            return self.round(x, f, training and self.stochastic, seed_gen)
+
+        x = wrap_sm_fn(x, k, i, f, training, quant_fn)
+        return x
+
+    def __call__(self, x, k, i, f, training=False, seed_gen=None):
+        i = ops.stop_gradient(ops.maximum(i, -f)) + (i - ops.stop_gradient(i))  # type: ignore
+        if self.stochastic and training:
+            assert seed_gen is not None, 'Seed generator must be provided for stochastic rounding.'
+        if self.overflow_mode != 'WRAP_SM':
+            return self.forward(x, k, i, f, training, seed_gen)
+        else:
+            return self.forward_wrap_sm(x, k, i, f, training, seed_gen)
+
+
+def _get_fixed_quantizer(round_mode: str = 'TRN', overflow_mode: str = 'WRAP'):
+    """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
+    """
+    quantizer = FixedPointQuantizer(round_mode, overflow_mode)
+    return quantizer  # type: ignore

{quantizers-1.1.0 → quantizers-1.2.0}/src/quantizers/fixed_point/fixed_point_ops.py

@@ -1,16 +1,13 @@
-from collections.abc import Callable
-from typing import TYPE_CHECKING, Any, TypeVar
+from typing import TYPE_CHECKING, TypeVar
 
 from numpy.typing import ArrayLike
 
 if TYPE_CHECKING:
-    from keras.api.random import SeedGenerator
+    pass
 T = TypeVar('T', bound=ArrayLike)
 
 
-def get_fixed_quantizer(
-    round_mode: str = 'TRN', overflow_mode: str = 'WRAP'
-) -> Callable[[T, Any, Any, Any, bool | None, 'SeedGenerator | None'], T]:
+def get_fixed_quantizer(round_mode: str = 'TRN', overflow_mode: str = 'WRAP'):
     """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.
 

{quantizers-1.1.0 → quantizers-1.2.0}/src/quantizers/fixed_point/fixed_point_ops_np.py

@@ -135,12 +135,8 @@ def get_fixed_quantizer_np(round_mode: str = 'TRN', overflow_mode: str = 'WRAP')
     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}')
-    sat_fn = saturation_mode_registry_np.get(overflow_mode, None)
-    if sat_fn is None:
-        raise ValueError(f'Unknown saturation mode: {overflow_mode}')
+    round_fn_scaled = round_mode_registry_scaled_np[round_mode]
+    sat_fn = saturation_mode_registry_np[overflow_mode]
 
     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.'

{quantizers-1.1.0 → quantizers-1.2.0}/src/quantizers/minifloat/float_point_ops.py

@@ -1,4 +1,7 @@
-def float_quantize(x, m, e, e0=0.0):
+from typing import Any
+
+
+def float_quantize(x, m, e, e0: Any = 0.0):
     """Quantize an array to floatlet (m mantissa bits, excl. sign bit, e exponent bits) format. Tentative gradient impl."""
     from ._float_point_ops import _float_quantize
 

{quantizers-1.1.0 → quantizers-1.2.0}/src/quantizers/stateful_quantizers.py

@@ -91,6 +91,9 @@ class FixedQ:
         if use_numpy:
             q = get_fixed_quantizer_np(self.rounding, self.overflow)
         else:
+            from keras.api import ops
+
+            k, i, f = ops.convert_to_tensor(k), ops.convert_to_tensor(i), ops.convert_to_tensor(f)
             q = get_fixed_quantizer(self.rounding, self.overflow)
         return q(x, k, i, f, False, None)  # type: ignore
 
@@ -128,6 +131,9 @@ class MinifloatQ:
         if use_numpy:
             return float_quantize_np(x, s, e, e0)
         else:
+            from keras.api import ops
+
+            s, e, e0 = ops.convert_to_tensor(s), ops.convert_to_tensor(e), ops.convert_to_tensor(e0)
             return float_quantize(x, s, e, e0)
 
     def __repr__(self):

{quantizers-1.1.0 → quantizers-1.2.0/src/quantizers.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.4
 Name: quantizers
-Version: 1.1.0
+Version: 1.2.0
 Author-email: Chang Sun <chsun@cern.ch>
 License: GNU Lesser General Public License v3 (LGPLv3)
 Project-URL: repository, https://github.com/calad0i/quantizers
@@ -24,6 +24,7 @@ Requires-Dist: pytest-env; extra == "test"
 Requires-Dist: keras>=3.0; extra == "test"
 Requires-Dist: cppyy; extra == "test"
 Requires-Dist: jax; extra == "test"
+Dynamic: license-file
 
 # Quantizers
 [![PyPI version](https://badge.fury.io/py/quantizers.svg)](https://badge.fury.io/py/quantizers)

{quantizers-1.1.0 → quantizers-1.2.0}/test/test_quantizer_match.py

@@ -18,7 +18,7 @@ def register_cpp():
 
 @pytest.fixture(scope='session')
 def data():
-    arr = np.random.randint(-2**15, 2**15, 1000) * 2.**-8
+    arr = np.random.randint(-(2**15), 2**15, 1000) * 2.0**-8
     return ops.array(arr.astype(np.float32))
 
 
@@ -53,19 +53,20 @@ def c_quantize_ternary(x):
     return r
 
 
-@pytest.mark.parametrize('fixed_round_mode', ['TRN', 'TRN_ZERO', 'RND', 'S_RND', 'S_RND_CONV', 'RND_CONV', 'RND_ZERO', 'RND_INF', 'RND_MIN_INF'])
+@pytest.mark.parametrize(
+    'fixed_round_mode', ['TRN', 'TRN_ZERO', 'RND', 'S_RND', 'S_RND_CONV', 'RND_CONV', 'RND_ZERO', 'RND_INF', 'RND_MIN_INF']
+)
 @pytest.mark.parametrize('fixed_overflow_mode', ['WRAP', 'WRAP_SM', 'SAT', 'SAT_SYM'])
 @pytest.mark.parametrize('k', [0, 1])
 @pytest.mark.parametrize('b', [2, 7, 12])
 @pytest.mark.parametrize('i', [-5, 0, 8])
 def test_fixed_quantizer_forward(fixed_round_mode, fixed_overflow_mode, k, b, i, data, register_cpp):
-
     k, i, f = k, i, b - i
     W, I = k + b, k + i
     if fixed_overflow_mode == 'WRAP_SM':
         if k == 0:
             pytest.skip('WRAP_SM does not support k=0')
-        if fixed_round_mode not in ('RND_CONV', 'RND'):
+        if fixed_round_mode not in ('RND_CONV', 'RND', 'S_RND', 'S_RND_CONV'):
             pytest.skip('WRAP_SM only supports RND-like rounding')
 
     fixed_q = FixedQ(W, I, k, fixed_round_mode, fixed_overflow_mode)
@@ -73,25 +74,31 @@ def test_fixed_quantizer_forward(fixed_round_mode, fixed_overflow_mode, k, b, i,
 
     arr_py_fixed_np = np.array(fixed_q(np.array(data)))
     arr_py_fixed = np.array(fixed_q(data))
+    mismatch = np.where(arr_py_fixed_np != arr_py_fixed)[0]
+    assert len(mismatch) == 0, f"""Fixed quantizer has inconsistent behavior with numpy implementation:
+        [* {len(mismatch)} mismatches, {min(len(mismatch), 5)} shown *]
+        np: {arr_py_fixed_np[mismatch][:5]}
+        keras: {arr_py_fixed[mismatch][:5]}
+        in: {data[mismatch][:5] * 2.0**8}
+    """
     assert np.all(arr_py_fixed_np == arr_py_fixed), 'numpy / keras implementation inconsistent'
 
     arr_c_fixed = c_quantize_fixed(data, k, i, f, fixed_round_mode, fixed_overflow_mode)
     arr_c_fixed_np = np.array([float(x) for x in arr_c_fixed])
 
     mismatch = np.where(arr_py_fixed != arr_c_fixed_np)[0]
-    assert len(mismatch) == 0, f'''Fixed quantizer has inconsistent behavior with C++ implementation:
+    assert len(mismatch) == 0, f"""Fixed quantizer has inconsistent behavior with C++ implementation:
         [* {len(mismatch)} mismatches, {min(len(mismatch), 5)} shown *]
         C++: {arr_c_fixed_np[mismatch][:5]}
         Py: {arr_py_fixed[mismatch][:5]}
         in: {data[mismatch][:5]}
-    '''
+    """
 
 
 @pytest.mark.parametrize('M', [2, 6])
 @pytest.mark.parametrize('E', [2, 6])
 @pytest.mark.parametrize('E0', [0, 3, 5])
 def test_minifloat_quantizer_forward(M, E, E0, data, register_cpp):
-
     q = MinifloatQ(M, E, E0)
     print(q)
     arr_py_float = np.array(q(data))
@@ -102,13 +109,13 @@ def test_minifloat_quantizer_forward(M, E, E0, data, register_cpp):
     arr_c_float_np = np.array([float(x) for x in arr_c_float])
 
     mismatch = np.where(arr_py_float != arr_c_float_np)[0]
-    assert len(mismatch) == 0, f'''Float quantizer has inconsistent behavior with C++ implementation:
+    assert len(mismatch) == 0, f"""Float quantizer has inconsistent behavior with C++ implementation:
         [* {len(mismatch)} mismatches, {min(len(mismatch), 5)} shown *]
         [* Up to 5 shown]
         C++: {arr_c_float_np[mismatch][:5]}
         Py: {arr_py_float[mismatch][:5]}
         in: {data[mismatch][:5]}
-    '''
+    """
 
 
 @pytest.mark.parametrize('fixed_round_mode', ['RND'])
@@ -120,7 +127,6 @@ def test_minifloat_quantizer_forward(M, E, E0, data, register_cpp):
 @pytest.mark.parametrize('E', [2, 4])
 @pytest.mark.parametrize('E0', [1, 8])
 def test_fixed_float_mult(fixed_round_mode, fixed_overflow_mode, k, i, f, M, E, E0, data, register_cpp):
-
     if fixed_overflow_mode == 'WRAP_SM':
         if k == 0:
             pytest.skip('WRAP_SM does not support k=0')
@@ -137,47 +143,47 @@ def test_fixed_float_mult(fixed_round_mode, fixed_overflow_mode, k, i, f, M, E,
 
     mismatch = np.where(mult_py != mult_c)[0]
     print(len(mismatch))
-    assert len(mismatch) == 0, f'''Multiplication has inconsistent behavior with C++ implementation:
+    assert len(mismatch) == 0, f"""Multiplication has inconsistent behavior with C++ implementation:
         [* {len(mismatch)} mismatches, {min(len(mismatch), 5)} shown *]
         C++: {mult_c[mismatch][:5]}
         Py: {mult_py[mismatch][:5]}
         fx in: {arr_c_fixed_np[mismatch][:5]}
         fp in: {arr_c_float_np[mismatch][:5]}
-    '''
+    """
 
 
 def test_binary(data, register_cpp):
     q = BinaryQ()
     print(q)
-    arr_py_binary = q(data)
-    arr_py_binary_np = np.array(q(np.array(data)))
+    arr_py_binary = np.array(q(data))
+    arr_py_binary_np = q(np.array(data))
     assert np.all(arr_py_binary_np == arr_py_binary), 'numpy / keras implementation inconsistent'
 
     arr_c_binary = c_quantize_binary(data)
     arr_c_binary_np = np.array([float(x) for x in arr_c_binary])
 
     mismatch = np.where(arr_py_binary != arr_c_binary_np)[0]
-    assert len(mismatch) == 0, f'''Binary quantizer has inconsistent behavior with C++ implementation:
+    assert len(mismatch) == 0, f"""Binary quantizer has inconsistent behavior with C++ implementation:
         [* {len(mismatch)} mismatches, {min(len(mismatch), 5)} shown *]
         C++: {arr_c_binary_np[mismatch][:5]}
         Py: {arr_py_binary[mismatch][:5]}
         in: {data[mismatch][:5]}
-    '''
+    """
 
 
 def test_ternary(data, register_cpp):
     q = TernaryQ()
     print(q)
-    arr_py_ternary = q(data)
-    arr_py_ternary_np = np.array(q(np.array(data)))
+    arr_py_ternary = np.array(q(data))
+    arr_py_ternary_np = q(np.array(data))
     assert np.all(arr_py_ternary_np == arr_py_ternary), 'numpy / keras implementation inconsistent'
 
     arr_c_ternary = c_quantize_ternary(data)
     arr_c_ternary_np = np.array([float(x) for x in arr_c_ternary])
     mismatch = np.where(arr_py_ternary != arr_c_ternary_np)[0]
-    assert len(mismatch) == 0, f'''Ternary quantizer has inconsistent behavior with C++ implementation:
+    assert len(mismatch) == 0, f"""Ternary quantizer has inconsistent behavior with C++ implementation:
         [* {len(mismatch)} mismatches, {min(len(mismatch), 5)} shown *]
         C++: {arr_c_ternary_np[mismatch][:5]}
         Py: {arr_py_ternary[mismatch][:5]}
         in: {data[mismatch][:5]}
-    '''
+    """

{quantizers-1.1.0 → quantizers-1.2.0}/test/test_stateless_quantizer.py

@@ -4,9 +4,16 @@ import jax
 import jax.numpy as jnp
 import numpy as np
 import pytest
+from keras.api.random import SeedGenerator
 
-from quantizers import binary_quantize, float_decompose, float_decompose_np, float_quantize, get_fixed_quantizer, ternary_quantize
-from quantizers.fixed_point._fixed_point_ops import round_mode_registry
+from quantizers import (
+    binary_quantize,
+    float_decompose,
+    float_decompose_np,
+    float_quantize,
+    get_fixed_quantizer,
+    ternary_quantize,
+)
 
 
 @pytest.fixture(scope='module')
@@ -30,26 +37,28 @@ def xxx_mee0():
     M = jax.random.randint(k1, N, 1, 9).astype(jnp.float32)
     E = jax.random.randint(k2, N, 1, 4).astype(jnp.float32)
     E0 = jax.random.randint(k3, N, -8, 8).astype(jnp.float32)
-    x = jax.random.uniform(k4, N, jnp.float32, -1., 1.)
+    x = jax.random.uniform(k4, N, jnp.float32, -1.0, 1.0)
 
     sign = jnp.sign(x)
     x = jnp.abs(x)
-    m_eps = 2.**-M
-    _max = (2 - m_eps) * 2.**(2.**(E - 1) - 1 + E0)
-    _min_pos_normal = 2.**(-2.**(E - 1) + E0 + 1)
-    _min_pos_subnormal = m_eps * 2.**(-2.**(E - 1) + E0)
+    m_eps = 2.0**-M
+    _max = (2 - m_eps) * 2.0 ** (2.0 ** (E - 1) - 1 + E0)
+    _min_pos_normal = 2.0 ** (-(2.0 ** (E - 1)) + E0 + 1)
+    _min_pos_subnormal = m_eps * 2.0 ** (-(2.0 ** (E - 1)) + E0)
     log_subnormal = jnp.log2(_min_pos_subnormal)
     log_normal = jnp.log2(_min_pos_normal)
     log_overflow = jnp.log2(_max)
 
-    x_normal = 2.**(x * (log_overflow - log_normal) + log_normal) * sign
-    x_subnormal = 2.**(x * (log_normal - log_subnormal) + log_subnormal - 1) * sign
-    x_overflow = 2.**(x * 2 + log_overflow + 1) * sign
+    x_normal = 2.0 ** (x * (log_overflow - log_normal) + log_normal) * sign
+    x_subnormal = 2.0 ** (x * (log_normal - log_subnormal) + log_subnormal - 1) * sign
+    x_overflow = 2.0 ** (x * 2 + log_overflow + 1) * sign
     xxx = (x_normal, x_subnormal, x_overflow)
     return xxx, M, E, E0
 
 
-@pytest.mark.parametrize('round_mode', ['TRN', 'TRN_ZERO', 'RND', 'S_RND', 'RND_CONV', 'S_RND_CONV', 'RND_ZERO', 'RND_INF', 'RND_MIN_INF'])
+@pytest.mark.parametrize(
+    'round_mode', ['TRN', 'TRN_ZERO', 'RND', 'S_RND', 'RND_CONV', 'S_RND_CONV', 'RND_ZERO', 'RND_INF', 'RND_MIN_INF']
+)
 @pytest.mark.parametrize('overflow_mode', ['WRAP', 'SAT', 'SAT_SYM', 'WRAP_SM'])
 def test_fixed_quantizer_grad(round_mode, overflow_mode, x_kif):
     if round_mode not in ('RND', 'RND_CONV') and overflow_mode == 'WRAP_SM':
@@ -57,44 +66,37 @@ def test_fixed_quantizer_grad(round_mode, overflow_mode, x_kif):
 
     quantizer = get_fixed_quantizer(round_mode, overflow_mode)
     assert callable(quantizer)
+    seed = SeedGenerator(42)
 
     x, k, i, f = x_kif
 
-    xq = quantizer(x, k, i, f, True)  # type: ignore
-    if not round_mode.startswith('S_') and 'WRAP' not in overflow_mode:
+    xq = quantizer(x, k, i, f, True, seed)  # type: ignore
+    if 'WRAP' not in overflow_mode and not round_mode.startswith('S_'):
         xq1 = quantizer(x, k, i, f, False)  # type: ignore
         assert jnp.all(xq == xq1)
 
     def abs_quantization_err(x, k, i, f):
-        xq = quantizer(x, k, i, f, True)  # type: ignore
+        xq = quantizer(x, k, i, f, True, seed)  # type: ignore
         err = jnp.abs(x - xq)
         return jnp.sum(err)
 
     dx, di, df = jax.grad(abs_quantization_err, (0, 2, 3))(x, 1, i, f)
 
     if overflow_mode == 'WRAP':
-        assert jnp.all(dx == 0), f'X grad Error'
-        assert jnp.all(di == 0), f'I grad Error'
-        assert jnp.all(df < 0), f'F grad Error'
+        assert jnp.all(dx == 0), 'X grad Error'
+        assert jnp.all(di == 0), 'I grad Error'
+        assert jnp.all(df < 0), 'F grad Error'
     elif overflow_mode in ('SAT', 'SAT_SYM'):
         for _dx in jnp.unique(dx):
-            assert _dx in (-1, 0, 1), f'X grad Error'
-        assert jnp.all(df <= 0), f'F grad Error'
-        assert jnp.all(di <= 0), f'I grad Error'
-        assert jnp.all((df < 0) | (di < 0) | (x == xq)), f'Grad Error for sat mode'
+            assert _dx in (-1, 0, 1), 'X grad Error'
+        assert jnp.all(df <= 0), 'F grad Error'
+        assert jnp.all(di <= 0), 'I grad Error'
+        assert jnp.all((df < 0) | (di < 0) | (x == xq)), 'Grad Error for sat mode'
     elif overflow_mode == 'WRAP_SM':
-        mask = jnp.abs(x) < 2.**i - 2.**f
+        mask = jnp.abs(x) < 2.0**i - 2.0**f
         assert jnp.all(df[mask] < 0), f'F grad Error {jnp.mean(df <= 0)}'
 
 
-@pytest.mark.parametrize('round_mode', ['TRN', 'TRN_ZERO', 'RND', 'RND_CONV', 'RND_ZERO', 'RND_INF', 'RND_MIN_INF'])
-def test_round_ste(round_mode, x_kif):
-    x = x_kif[0]
-    rnd_fn = round_mode_registry[round_mode]
-    grad_x = jax.grad(lambda x: jnp.sum(rnd_fn(x, True)))(x)
-    assert jnp.all(grad_x == 1), f'Gradient Error for {round_mode}'
-
-
 def test_float_quantizer_grad(xxx_mee0):
     xxx, M, E, E0 = xxx_mee0
     x_normal, x_subnormal, x_overflow = xxx
@@ -107,72 +109,71 @@ def test_float_quantizer_grad(xxx_mee0):
     dx, dm, de, de0 = jax.grad(abs_quantization_err, range(4))(x_normal, M, E, E0)
     xq = float_quantize(x_normal, M, E, E0)
     mask = x_normal != xq
-    assert jnp.all(dx == 0), f'Normal Number X grad Error'
+    assert jnp.all(dx == 0), 'Normal Number X grad Error'
     assert jnp.all(dm[mask] < 0), f'Normal Number M grad Error: max={jnp.max(dm[mask])}>0'
-    assert jnp.all(de == 0), f'Normal Number E grad Error'
-    assert jnp.all(de0 == 0), f'Normal Number E0 grad Error'
+    assert jnp.all(de == 0), 'Normal Number E grad Error'
+    assert jnp.all(de0 == 0), 'Normal Number E0 grad Error'
 
     dx, dm, de, de0 = jax.grad(abs_quantization_err, range(4))(x_subnormal, M, E, E0)
     xq = float_quantize(x_subnormal, M, E, E0)
     mask = x_subnormal != xq
-    assert jnp.all(dx == 0), f'Subnormal Number X grad Error'
+    assert jnp.all(dx == 0), 'Subnormal Number X grad Error'
     assert jnp.all(dm[mask] < 0), f'Subnormal Number M grad Error: max={jnp.max(dm[mask])}>0'
     assert jnp.all(de[mask] < 0), f'Subnormal Number E grad Error: max={jnp.max(de[mask])}>0'
     assert jnp.all(de0[mask] > 0), f'Subnormal Number E0 grad Error: max={jnp.max(de0[mask])}>0'
 
     dx, dm, de, de0 = jax.grad(abs_quantization_err, range(4))(x_overflow, M, E, E0)
-    assert jnp.all(dx == 0), f'Overflow Number X grad Error'
-    assert jnp.all(dm == 0), f'Overflow Number M grad Error'
+    assert jnp.all(dx == 0), 'Overflow Number X grad Error'
+    assert jnp.all(dm == 0), 'Overflow Number M grad Error'
     assert jnp.all(de < 0), f'Overflow Number E grad Error: max={jnp.max(de)}>0'
     assert jnp.all(de0 < 0), f'Overflow Number E0 grad Error: max={jnp.max(de0)}>0'
 
 
 def test_float_decompose(xxx_mee0):
-
     xxx, M, E, E0 = xxx_mee0
     x_normal, x_subnormal, x_overflow = xxx
 
     mm, ee = float_decompose(x_normal, M, E, E0)
     mm_, ee_ = float_decompose_np(x_normal, M, E, E0)
-    assert np.all(mm == mm_) and np.all(ee == ee_), f'Float Decompose Error @ Normal Number'
+    assert np.all(mm == mm_) and np.all(ee == ee_), 'Float Decompose Error @ Normal Number'
 
-    xq_ = mm * 2.**ee  # type: ignore
+    xq_ = mm * 2.0**ee  # type: ignore
     xq = float_quantize(x_normal, M, E, E0)
-    assert jnp.all(xq == xq_), f'Float Decompose Error @ Normal Number'
-    assert jnp.all(jnp.abs(mm) < 2), f'Mantissa Error @ Normal Number'  # type: ignore
-    assert jnp.all(jnp.abs(mm) >= 1), f'Mantissa Error @ Normal Number'  # type: ignore
+    assert jnp.all(xq == xq_), 'Float Decompose Error @ Normal Number'
+    assert jnp.all(jnp.abs(mm) < 2), 'Mantissa Error @ Normal Number'  # type: ignore
+    assert jnp.all(jnp.abs(mm) >= 1), 'Mantissa Error @ Normal Number'  # type: ignore
 
     mm, ee = float_decompose(x_subnormal, M, E, E0)
     mm_, ee_ = float_decompose_np(x_subnormal, M, E, E0)
-    assert np.all(mm == mm_) and np.all(ee == ee_), f'Float Decompose Error @ Subnormal Number'
+    assert np.all(mm == mm_) and np.all(ee == ee_), 'Float Decompose Error @ Subnormal Number'
 
-    xq_ = mm * 2.**ee  # type: ignore
+    xq_ = mm * 2.0**ee  # type: ignore
     xq = float_quantize(x_subnormal, M, E, E0)
-    assert jnp.all(xq == xq_), f'Float Decompose Error @ Subnormal Number'
-    assert jnp.all(jnp.abs(mm) < 1), f'Mantissa Error @ Subnormal Number'  # type: ignore
+    assert jnp.all(xq == xq_), 'Float Decompose Error @ Subnormal Number'
+    assert jnp.all(jnp.abs(mm) < 1), 'Mantissa Error @ Subnormal Number'  # type: ignore
 
     mm, ee = float_decompose(x_overflow, M, E, E0)
     mm_, ee_ = float_decompose_np(x_overflow, M, E, E0)
-    assert np.all(mm == mm_) and np.all(ee == ee_), f'Float Decompose Error'
+    assert np.all(mm == mm_) and np.all(ee == ee_), 'Float Decompose Error'
 
-    xq_ = mm * 2.**ee  # type: ignore
+    xq_ = mm * 2.0**ee  # type: ignore
     xq = float_quantize(x_overflow, M, E, E0)
-    assert jnp.all(xq == xq_), f'Float Decompose Error'
-    assert jnp.all(jnp.abs(mm) < 2), f'Mantissa Error @ Overflow Number'  # type: ignore
-    assert jnp.all(jnp.abs(mm) >= 1), f'Mantissa Error @ Overflow Number'  # type: ignore
+    assert jnp.all(xq == xq_), 'Float Decompose Error'
+    assert jnp.all(jnp.abs(mm) < 2), 'Mantissa Error @ Overflow Number'  # type: ignore
+    assert jnp.all(jnp.abs(mm) >= 1), 'Mantissa Error @ Overflow Number'  # type: ignore
 
 
 def test_binary_quantizer_grad(x_kif):
     x, *_ = x_kif
     xq = binary_quantize(x)
-    assert jnp.all((xq == -1) | (xq == 1)), f'Binary Quantizer Error'
+    assert jnp.all((xq == -1) | (xq == 1)), 'Binary Quantizer Error'
     grad = jax.grad(lambda x: jnp.sum(binary_quantize(x)))(x)
-    assert jnp.all(grad > 0), f'Gradient Error'
+    assert jnp.all(grad > 0), 'Gradient Error'
 
 
 def test_ternary_quantizer_grad(x_kif):
     x, *_ = x_kif
     xq = ternary_quantize(x)
-    assert jnp.all((xq == -1) | (xq == 0) | (xq == 1)), f'Ternary Quantizer Error'
+    assert jnp.all((xq == -1) | (xq == 0) | (xq == 1)), 'Ternary Quantizer Error'
     grad = jax.grad(lambda x: jnp.sum(ternary_quantize(x)))(x)
-    assert jnp.all(grad > 0), f'Gradient Error'
+    assert jnp.all(grad > 0), 'Gradient Error'

quantizers-1.1.0/src/quantizers/fixed_point/_fixed_point_ops.py

@@ -1,289 +0,0 @@
-from collections.abc import Callable
-from typing import Any, TypeVar
-
-import keras
-from keras import ops
-from keras.api.random import SeedGenerator
-from numpy.typing import ArrayLike
-
-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]] = {}
-
-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.0)
-        dmax = ops.where(mask1, dy, 0.0)
-        dmin = ops.where(mask2, dy, 0.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.0**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.0) * 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.0 ** (bk - 1)
-    return (xs + bias) % (2.0**bk) - bias
-
-
-@sat_mode('SAT')
-def sat(x, k, i, f, training=None):
-    f_eps = 2.0 ** (-f)
-    __max = 2.0**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.0 ** (-f)
-    _max = 2.0**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.0**-f
-    high = 2.0**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}')
-    sat_fn = saturation_mode_registry.get(overflow_mode, None)
-    if sat_fn is None:
-        raise ValueError(f'Unknown saturation mode: {overflow_mode}')
-
-    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.0**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.0) * 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()