quantizers 1.0.1__tar.gz → 1.2.0__tar.gz

{quantizers-1.0.1 → quantizers-1.2.0}/.pre-commit-config.yaml

@@ -1,15 +1,12 @@
 exclude: '\.svg$|ap_types/|_version.py'
 
 repos:
-- repo: https://github.com/hhatto/autopep8
-  rev: v2.3.1
+- repo: https://github.com/astral-sh/ruff-pre-commit
+  rev: v0.8.2
   hooks:
-  - id: autopep8
-    language_version: python3
-    args: ['--max-line-length=65535',
-           '--ignore=E731,E402',
-           '--in-place',
-           ]
+    - id: ruff
+      args: [--fix]
+    - id: ruff-format
 
 - repo: https://github.com/pre-commit/pre-commit-hooks
   rev: v4.6.0
@@ -36,12 +33,6 @@ repos:
   - id: pyupgrade
     args: ["--py310-plus"]
 
-- repo: https://github.com/PyCQA/isort
-  rev: 5.13.2
-  hooks:
-  - id: isort
-    args: ['--line-length=65535']
-
 - repo: https://github.com/pre-commit/mirrors-clang-format
   rev: v18.1.8
   hooks:

{quantizers-1.0.1/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.0.1
+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.0.1 → quantizers-1.2.0}/pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "quantizers"
-version = "1.0.1"
+version = "1.2.0"
 authors = [{ name = "Chang Sun", email = "chsun@cern.ch" }]
 description = ""
 readme = "README.md"
@@ -70,3 +70,28 @@ addopts = "--import-mode=importlib --cov=src/quantizers --cov-report=xml --cov-r
 
 [tool.coverage.run]
 omit = ["*/_version.py"]
+
+[tool.coverage.report]
+exclude_lines = ['pragma: no cover', 'if TYPE_CHECKING:']
+
+[tool.ruff]
+include = ["src/**/*.py"]
+
+line-length = 130
+indent-width = 4
+
+target-version = "py310"
+
+[tool.ruff.format]
+quote-style = "single"
+indent-style = "space"
+skip-magic-trailing-comma = false
+docstring-code-format = true
+docstring-code-line-length = 130
+
+[tool.ruff.lint]
+select = ["E", "W", "F", "I", "F401"]
+ignore = ["E501", "F403", "F405", "E741"]
+explicit-preview-rules = true
+fixable = ["ALL"]
+unfixable = []

{quantizers-1.0.1 → quantizers-1.2.0}/src/quantizers/__init__.py

@@ -2,3 +2,5 @@ from .binary import *
 from .fixed_point import *
 from .minifloat import *
 from .stateful_quantizers import BinaryQ, FixedQ, MinifloatQ, TernaryQ
+
+__all__ = ['BinaryQ', 'FixedQ', 'MinifloatQ', 'TernaryQ']

{quantizers-1.0.1 → 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.0.1'
-__version_tuple__ = version_tuple = (1, 0, 1)
+__version__ = version = '1.2.0'
+__version_tuple__ = version_tuple = (1, 2, 0)

quantizers-1.2.0/src/quantizers/binary/__init__.py

@@ -0,0 +1,4 @@
+from .binary_ops import binary_quantize, ternary_quantize
+from .binary_ops_np import binary_quantize_np, ternary_quantize_np
+
+__all__ = ['binary_quantize', 'ternary_quantize', 'binary_quantize_np', 'ternary_quantize_np']

quantizers-1.0.1/src/quantizers/binary/binary_ops.py → quantizers-1.2.0/src/quantizers/binary/_binary_ops.py

@@ -1,28 +1,27 @@
 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 _binary_quantize(x):
+    r = ops.where(x > 0, 1.0, -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 dy * (1.0 - 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 _ternary_quantize(x):
+    r = ops.where(x > 0.5, 1.0, ops.where(x < -0.5, -1.0, 0.0))
 
     def grad(*args, upstream=None):
         if upstream is None:
             (upstream,) = args
         dy = upstream
-        return dy * (1. - ops.tanh(x)**2)  # type: ignore
+        return dy * (1.0 - ops.tanh(x) ** 2)  # type: ignore
 
     return ops.stop_gradient(r), grad

quantizers-1.2.0/src/quantizers/binary/binary_ops.py

@@ -0,0 +1,10 @@
+def binary_quantize(x):
+    from ._binary_ops import _binary_quantize
+
+    return _binary_quantize(x)
+
+
+def ternary_quantize(x):
+    from ._binary_ops import _ternary_quantize
+
+    return _ternary_quantize(x)

quantizers-1.2.0/src/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.0, -1.0)
+
+
+def ternary_quantize_np(x):
+    return np.where(x > 0.5, 1.0, np.where(x < -0.5, -1.0, 0.0))

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

@@ -0,0 +1,4 @@
+from .fixed_point_ops import get_fixed_quantizer
+from .fixed_point_ops_np import get_fixed_quantizer_np
+
+__all__ = ['get_fixed_quantizer', 'get_fixed_quantizer_np']

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.2.0/src/quantizers/fixed_point/fixed_point_ops.py

@@ -0,0 +1,19 @@
+from typing import TYPE_CHECKING, TypeVar
+
+from numpy.typing import ArrayLike
+
+if TYPE_CHECKING:
+    pass
+T = TypeVar('T', bound=ArrayLike)
+
+
+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
+    """
+    from ._fixed_point_ops import _get_fixed_quantizer
+
+    return _get_fixed_quantizer(round_mode, overflow_mode)

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

@@ -2,7 +2,6 @@ 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]] = {}
@@ -16,7 +15,7 @@ def rnd_mode_np(names: str | list[str] | tuple[str, ...]):
 
     def inner(func):
         def wrapper(x, f, training=None, seed_gen=None):
-            scale = 2.**f
+            scale = 2.0**f
             sx = x * scale
             sxq = func(sx)
             xq = sxq / scale
@@ -33,6 +32,7 @@ def rnd_mode_np(names: str | list[str] | tuple[str, ...]):
         round_mode_registry_np[func.__name__.upper()] = func
 
         return func
+
     return inner
 
 
@@ -83,6 +83,7 @@ def sat_mode_np(name: str | list | tuple):
             saturation_mode_registry_np[name] = func
         saturation_mode_registry_np[func.__name__.upper()] = func
         return func
+
     return inner
 
 
@@ -90,14 +91,14 @@ def sat_mode_np(name: str | list | tuple):
 def wrap(x, k, i, f, training=None):
     xs = x
     bk = i + k
-    bias = k * 2.**(bk - 1)
-    return ((xs + bias) % (2.**bk) - bias)
+    bias = k * 2.0 ** (bk - 1)
+    return (xs + bias) % (2.0**bk) - bias
 
 
 @sat_mode_np('SAT')
 def sat(x, k, i, f, training=None):
-    f_eps = 2.**(-f)
-    __max = 2.**i
+    f_eps = 2.0 ** (-f)
+    __max = 2.0**i
     _max = __max - f_eps
     _min = -__max * k
     return np.clip(x, _min, _max)
@@ -105,16 +106,16 @@ def sat(x, k, i, f, training=None):
 
 @sat_mode_np('SAT_SYM')
 def sat_sym(x, k, i, f, training=None):
-    f_eps = 2.**(-f)
-    _max = 2.**i - f_eps
+    f_eps = 2.0 ** (-f)
+    _max = 2.0**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
+    eps = 2.0**-f
+    high = 2.0**i - eps
     low = -(high + eps) * k
     interval = 2 ** (i + k)
     c1 = ((x) / interval) % 2 >= 1
@@ -134,24 +135,22 @@ 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}")  # 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
+    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."
+        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."
+        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':
@@ -160,4 +159,5 @@ def get_fixed_quantizer_np(round_mode: str = 'TRN', overflow_mode: str = 'WRAP')
         if overflow_mode == 'WRAP':
             x = sat_fn(x, k, _i, f, training)
         return x
+
     return quantizer

quantizers-1.2.0/src/quantizers/minifloat/__init__.py

@@ -0,0 +1,4 @@
+from .float_point_ops import float_decompose, float_quantize
+from .float_point_ops_np import float_decompose_np, float_quantize_np
+
+__all__ = ['float_quantize', 'float_decompose', 'float_quantize_np', 'float_decompose_np']

quantizers-1.0.1/src/quantizers/minifloat/float_point_ops.py → quantizers-1.2.0/src/quantizers/minifloat/_float_point_ops.py

@@ -1,21 +1,20 @@
-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.'''
+def _float_quantize(x, m, e, e0=0.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_high = ops.maximum(2.0 ** (e - 1), 1.0)
     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)
+    scale = 2.0 ** (e_act - m + 1)
     sig = x / scale
     qsig = ops.round(sig)
-    _sig_high = 2.**m
+    _sig_high = 2.0**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
@@ -26,31 +25,32 @@ def float_quantize(x, m, e, e0=0.):
         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
+        dm = ops.where(mask_m, scale * (sig - qsig) * dy * ops.log(2.0), 0.0)
+        d_exp = ops.where(mask_e, (x - qx) * dy * ops.log(2.0), 0.0)
+        de = d_exp * ops.log(2.0) * _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.'''
+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_high = 2.0 ** (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)
+    scale = 2.0 ** (e_act - m + 1)
     sig = x / scale
     qsig = ops.round(sig)
-    cond = (e_act != e_high) & (ops.abs(qsig) == 2.**m)
+    cond = (e_act != e_high) & (ops.abs(qsig) == 2.0**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)
+    qsig = ops.where(cond, qsig / 2.0, qsig)  # type: ignore
+    r_mantissa_bound = 2 - 2.0 ** (-m + 1)
+    r_mantissa = ops.clip(qsig * 2.0 ** (-m + 1), -r_mantissa_bound, r_mantissa_bound)
     r_exponent = e_act
     return r_mantissa, r_exponent

quantizers-1.2.0/src/quantizers/minifloat/float_point_ops.py

@@ -0,0 +1,15 @@
+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
+
+    return _float_quantize(x, m, e, e0)
+
+
+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."""
+    from ._float_point_ops import _float_decompose
+
+    return _float_decompose(x, m, e, e0)