Skip to main content
API REFERENCE/omnibias-binary

omnibias.binary.torch.ops.quantize

0.1.0a2AlphaPython >=3.10Apache-2.0

Binary and multilevel quantization with explicit hard-forward and surrogate-gradient contracts.

import omnibias.binary.torch.ops.quantize
Module contract & mathematical conventions

Torch closed-form quantization-gradient kernels.

Shared formulas do not imply universal floating-point identity. Verify dtype, shape, derivative-order support, and the stated numerical assumptions for the path you use.

Module source ↗Source revision 9a2d73df0573Running reference examples →
FUNCTION

binarize​

omnibias.binary.torch.ops.binarizeSource ↗
binarize(z: Tensor, beta: float=10.0) -> Tensor

Hard sign(z) in {-1, +1} (sign(0)=+1); Riccati surrogate backward.

Use a hard two-level decision in the forward pass while supplying an analytic training surrogate in reverse mode. Nonnegative inputs, including zero, map to +1; negative inputs map to -1. The cotangent with respect to z uses beta * (1 - tanh(beta*z)**2), not the classical derivative of sign. The result keeps the input shape. Choose a positive finite scalar beta: large magnitudes saturate the surrogate away from zero and concentrate its gradient near the boundary. The hard outputs themselves do not depend on beta. PyTorch uses a custom autograd.Function for hard quantizers.

Parameters

NameType & defaultContract
zpositional-or-keywordTensorRequiredFloating-point preactivation array of any shape; output has the same shape.
betapositional-or-keywordfloatDefault: 10.0Scalar inverse temperature of the backward surrogate. It does not change the hard forward threshold.
ReturnsTensor
Complete runnable example · binary torch

This is the complete downloadable script, tested against the pinned installed distribution. It includes initialization and assertions.

Download complete script ↓
"""Hard quantization, analytic surrogate derivatives, and an annealing schedule.

python -m pip install --pre "omnibias-binary==0.1.0a2" "omnibias-torch==0.5.0rc1"
"""
import math
import torch
from omnibias.binary import BetaAnnealScheduler
from omnibias.binary.torch.ops import (
    binarize, binarize01, heaviside, ternarize, kbit_quantize,
    riccati_tanh_derivative, riccati_sigmoid_derivative,
    surrogate_tower, surrogate_jet, curvature_corrected_slope, binarize_curvature,
)

z = torch.tensor([-0.75, -0.25, 0.0, 0.25, 0.75], dtype=torch.float64, requires_grad=True)
beta = 2.0
hard = binarize(z, beta)
torch.testing.assert_close(hard, torch.tensor([-1., -1., 1., 1., 1.], dtype=z.dtype))
# This is a specified surrogate gradient, not the classical derivative of sign.
(grad,) = torch.autograd.grad(hard.sum(), z)
torch.testing.assert_close(grad, beta * (1 - torch.tanh(beta*z).square()))
step = heaviside(z, beta)
torch.testing.assert_close(step, binarize01(z, beta))
(step_grad,) = torch.autograd.grad(step.sum(), z)
s = torch.sigmoid(beta*z)
torch.testing.assert_close(step_grad, beta*s*(1-s))
torch.testing.assert_close(ternarize(z, beta, delta=0.5), torch.tensor([-1., 0., 0., 0., 1.], dtype=z.dtype))
quantized = kbit_quantize(z, bits=2, lo=-1.0, hi=1.0, beta=beta)
assert quantized.shape == z.shape and torch.isfinite(quantized).all()
assert quantized.min() >= -1 and quantized.max() <= 1
(quantized_grad,) = torch.autograd.grad(quantized.sum(), z)
assert torch.isfinite(quantized_grad).all() and quantized_grad.abs().sum() > 0

t = torch.tanh(z)
torch.testing.assert_close(riccati_tanh_derivative(t, 1), 1-t.square())
s = torch.sigmoid(z)
torch.testing.assert_close(riccati_sigmoid_derivative(s, 1), s*(1-s))
tower = surrogate_tower(z, beta, order=3)
jet = surrogate_jet(z, beta, order=3)
assert tower.shape == jet.shape == (4, 5)
for order in range(4):
    torch.testing.assert_close(jet[order]*math.factorial(order), tower[order])
torch.testing.assert_close(curvature_corrected_slope(z, beta, window=0.0), tower[1])
(curvature_grad,) = torch.autograd.grad(binarize_curvature(z, beta).sum(), z)
torch.testing.assert_close(curvature_grad, curvature_corrected_slope(z, beta))

schedule = BetaAnnealScheduler(1.0, 4.0, num_steps=2, schedule="linear")
assert [schedule.step() for _ in range(3)] == [1.0, 2.5, 4.0]
assert schedule.current_step == 3 and schedule.value(100) == 4.0
schedule.reset()
assert schedule.current_step == 0 and schedule.step() == 1.0
print("quantizers: (5,); surrogate tower/jet: (4, 5); gradients and schedule passed")
python binary_torch.py
Executed regression · shape preservation

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗
def test_shape_preservation() -> None:
    z = torch.as_tensor(Z_2D, dtype=torch.float64)
    assert tq.binarize(z).shape == Z_2D.shape
    assert tq.ternarize(z).shape == Z_2D.shape
    assert tq.kbit_quantize(z, bits=2).shape == Z_2D.shape
python -m pytest "packages/omnibias-binary/tests/test_quantize.py::test_shape_preservation" -q
Read the implementation
def binarize(z: Tensor, beta: float = 10.0) -> Tensor:
    """Hard ``sign(z)`` in ``{-1, +1}`` (``sign(0)=+1``); Riccati surrogate backward."""
    return _BinarizeFn.apply(z, beta)
Other public import paths
  • omnibias.binary.torch.ops.binarize
  • omnibias.binary.torch.ops.quantize.binarize
FUNCTION

binarize01​

omnibias.binary.torch.ops.binarize01Source ↗
binarize01(z: Tensor, beta: float=10.0) -> Tensor

Hard Heaviside step in {0, 1} (H(0)=1); sigmoid(beta z) surrogate backward.

The {0, 1} codomain twin of binarize, using the Eulerian sigmoid_polynomial_coeffs Riccati tower instead of the tanh one. It is affinely conjugate to binarize: binarize01(z, beta) equals (binarize(z, beta / 2) + 1) / 2 in both forward and backward.

Choose this quantizer when the downstream representation requires bits rather than signed spins. It returns 1 at nonnegative z, including the threshold, and 0 below it. Reverse-mode differentiation uses beta * sigmoid(beta*z) * (1 - sigmoid(beta*z)). This is a deliberately supplied smooth surrogate; the discontinuous hard step has no classical derivative at zero. Shape is preserved. Positive finite beta sets the width of the gradient-bearing region; saturation can suppress training signals. PyTorch uses a custom autograd.Function for hard quantizers.

Parameters

NameType & defaultContract
zpositional-or-keywordTensorRequiredFloating-point preactivations of any shape.
betapositional-or-keywordfloatDefault: 10.0Scalar sharpness for the sigmoid backward surrogate; the forward outputs remain in {0, 1}.
ReturnsTensor
Complete runnable example · binary torch

This is the complete downloadable script, tested against the pinned installed distribution. It includes initialization and assertions.

Download complete script ↓
"""Hard quantization, analytic surrogate derivatives, and an annealing schedule.

python -m pip install --pre "omnibias-binary==0.1.0a2" "omnibias-torch==0.5.0rc1"
"""
import math
import torch
from omnibias.binary import BetaAnnealScheduler
from omnibias.binary.torch.ops import (
    binarize, binarize01, heaviside, ternarize, kbit_quantize,
    riccati_tanh_derivative, riccati_sigmoid_derivative,
    surrogate_tower, surrogate_jet, curvature_corrected_slope, binarize_curvature,
)

z = torch.tensor([-0.75, -0.25, 0.0, 0.25, 0.75], dtype=torch.float64, requires_grad=True)
beta = 2.0
hard = binarize(z, beta)
torch.testing.assert_close(hard, torch.tensor([-1., -1., 1., 1., 1.], dtype=z.dtype))
# This is a specified surrogate gradient, not the classical derivative of sign.
(grad,) = torch.autograd.grad(hard.sum(), z)
torch.testing.assert_close(grad, beta * (1 - torch.tanh(beta*z).square()))
step = heaviside(z, beta)
torch.testing.assert_close(step, binarize01(z, beta))
(step_grad,) = torch.autograd.grad(step.sum(), z)
s = torch.sigmoid(beta*z)
torch.testing.assert_close(step_grad, beta*s*(1-s))
torch.testing.assert_close(ternarize(z, beta, delta=0.5), torch.tensor([-1., 0., 0., 0., 1.], dtype=z.dtype))
quantized = kbit_quantize(z, bits=2, lo=-1.0, hi=1.0, beta=beta)
assert quantized.shape == z.shape and torch.isfinite(quantized).all()
assert quantized.min() >= -1 and quantized.max() <= 1
(quantized_grad,) = torch.autograd.grad(quantized.sum(), z)
assert torch.isfinite(quantized_grad).all() and quantized_grad.abs().sum() > 0

t = torch.tanh(z)
torch.testing.assert_close(riccati_tanh_derivative(t, 1), 1-t.square())
s = torch.sigmoid(z)
torch.testing.assert_close(riccati_sigmoid_derivative(s, 1), s*(1-s))
tower = surrogate_tower(z, beta, order=3)
jet = surrogate_jet(z, beta, order=3)
assert tower.shape == jet.shape == (4, 5)
for order in range(4):
    torch.testing.assert_close(jet[order]*math.factorial(order), tower[order])
torch.testing.assert_close(curvature_corrected_slope(z, beta, window=0.0), tower[1])
(curvature_grad,) = torch.autograd.grad(binarize_curvature(z, beta).sum(), z)
torch.testing.assert_close(curvature_grad, curvature_corrected_slope(z, beta))

schedule = BetaAnnealScheduler(1.0, 4.0, num_steps=2, schedule="linear")
assert [schedule.step() for _ in range(3)] == [1.0, 2.5, 4.0]
assert schedule.current_step == 3 and schedule.value(100) == 4.0
schedule.reset()
assert schedule.current_step == 0 and schedule.step() == 1.0
print("quantizers: (5,); surrogate tower/jet: (4, 5); gradients and schedule passed")
python binary_torch.py
Executed regression · forward is zero one

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗
def test_forward_is_zero_one() -> None:
    out = _np(tq.binarize01(torch.as_tensor(Z), beta=BETA))
    assert np.array_equal(out, np.where(Z >= 0, 1.0, 0.0))
    assert out[Z == 0.0][0] == 1.0  # H(0) = 1
python -m pytest "packages/omnibias-binary/tests/test_quantize01.py::test_forward_is_zero_one" -q
Read the implementation
def binarize01(z: Tensor, beta: float = 10.0) -> Tensor:
    """Hard Heaviside step in ``{0, 1}`` (``H(0)=1``); ``sigmoid(beta z)`` surrogate backward.

    The ``{0, 1}`` codomain twin of :func:`binarize`, using the Eulerian
    ``sigmoid_polynomial_coeffs`` Riccati tower instead of the ``tanh`` one. It is
    affinely conjugate to :func:`binarize`: ``binarize01(z, beta)`` equals
    ``(binarize(z, beta / 2) + 1) / 2`` in both forward and backward.
    """
    return _Binarize01Fn.apply(z, beta)
Other public import paths
  • omnibias.binary.torch.ops.binarize01
  • omnibias.binary.torch.ops.quantize.binarize01
FUNCTION

heaviside​

omnibias.binary.torch.ops.heavisideSource ↗
heaviside(z: Tensor, beta: float=10.0) -> Tensor

Alias of binarize01 (the hard {0, 1} Heaviside step).

Use the Heaviside spelling for the same hard bit-valued operation as binarize01. This is a callable wrapper with the same threshold convention H(0)=1 and the same sigmoid surrogate cotangent. For every input element, the forward result is 0 or 1 and has the original array shape. It should not be substituted for a smooth sigmoid when forward continuity is required. The downloadable example asserts both forward equality with binarize01 and the analytic reverse-mode gradient. PyTorch uses a custom autograd.Function for hard quantizers.

Parameters

NameType & defaultContract
zpositional-or-keywordTensorRequiredFloating-point preactivations; zero belongs to the positive side.
betapositional-or-keywordfloatDefault: 10.0Positive scalar sharpness used only by the supplied smooth backward rule.
ReturnsTensor
Complete runnable example · binary torch

This is the complete downloadable script, tested against the pinned installed distribution. It includes initialization and assertions.

Download complete script ↓
"""Hard quantization, analytic surrogate derivatives, and an annealing schedule.

python -m pip install --pre "omnibias-binary==0.1.0a2" "omnibias-torch==0.5.0rc1"
"""
import math
import torch
from omnibias.binary import BetaAnnealScheduler
from omnibias.binary.torch.ops import (
    binarize, binarize01, heaviside, ternarize, kbit_quantize,
    riccati_tanh_derivative, riccati_sigmoid_derivative,
    surrogate_tower, surrogate_jet, curvature_corrected_slope, binarize_curvature,
)

z = torch.tensor([-0.75, -0.25, 0.0, 0.25, 0.75], dtype=torch.float64, requires_grad=True)
beta = 2.0
hard = binarize(z, beta)
torch.testing.assert_close(hard, torch.tensor([-1., -1., 1., 1., 1.], dtype=z.dtype))
# This is a specified surrogate gradient, not the classical derivative of sign.
(grad,) = torch.autograd.grad(hard.sum(), z)
torch.testing.assert_close(grad, beta * (1 - torch.tanh(beta*z).square()))
step = heaviside(z, beta)
torch.testing.assert_close(step, binarize01(z, beta))
(step_grad,) = torch.autograd.grad(step.sum(), z)
s = torch.sigmoid(beta*z)
torch.testing.assert_close(step_grad, beta*s*(1-s))
torch.testing.assert_close(ternarize(z, beta, delta=0.5), torch.tensor([-1., 0., 0., 0., 1.], dtype=z.dtype))
quantized = kbit_quantize(z, bits=2, lo=-1.0, hi=1.0, beta=beta)
assert quantized.shape == z.shape and torch.isfinite(quantized).all()
assert quantized.min() >= -1 and quantized.max() <= 1
(quantized_grad,) = torch.autograd.grad(quantized.sum(), z)
assert torch.isfinite(quantized_grad).all() and quantized_grad.abs().sum() > 0

t = torch.tanh(z)
torch.testing.assert_close(riccati_tanh_derivative(t, 1), 1-t.square())
s = torch.sigmoid(z)
torch.testing.assert_close(riccati_sigmoid_derivative(s, 1), s*(1-s))
tower = surrogate_tower(z, beta, order=3)
jet = surrogate_jet(z, beta, order=3)
assert tower.shape == jet.shape == (4, 5)
for order in range(4):
    torch.testing.assert_close(jet[order]*math.factorial(order), tower[order])
torch.testing.assert_close(curvature_corrected_slope(z, beta, window=0.0), tower[1])
(curvature_grad,) = torch.autograd.grad(binarize_curvature(z, beta).sum(), z)
torch.testing.assert_close(curvature_grad, curvature_corrected_slope(z, beta))

schedule = BetaAnnealScheduler(1.0, 4.0, num_steps=2, schedule="linear")
assert [schedule.step() for _ in range(3)] == [1.0, 2.5, 4.0]
assert schedule.current_step == 3 and schedule.value(100) == 4.0
schedule.reset()
assert schedule.current_step == 0 and schedule.step() == 1.0
print("quantizers: (5,); surrogate tower/jet: (4, 5); gradients and schedule passed")
python binary_torch.py
Read the implementation
def heaviside(z: Tensor, beta: float = 10.0) -> Tensor:
    """Alias of :func:`binarize01` (the hard ``{0, 1}`` Heaviside step)."""
    return _Binarize01Fn.apply(z, beta)
Other public import paths
  • omnibias.binary.torch.ops.heaviside
  • omnibias.binary.torch.ops.quantize.heaviside
FUNCTION

kbit_quantize​

omnibias.binary.torch.ops.kbit_quantizeSource ↗
kbit_quantize(z: Tensor, bits: int=2, lo: float=-1.0, hi: float=1.0, beta: float=10.0) -> Tensor

Uniform k-bit quantize on [lo, hi]; tanh-step surrogate backward.

Quantize elementwise to 2**bits uniformly spaced levels from lo through hi. The forward path rounds (z-lo)/level_step and clips the level index, so values outside the interval saturate at its endpoints. Reverse mode instead sums shifted tanh surrogate slopes over the implementation’s level thresholds. That cotangent is not a numerical derivative of the rounded output. Use an integer bits >= 1 and finite hi > lo; the implementation explicitly rejects bits < 1, but does not validate every interval error. Both the surrogate workload and its threshold count grow exponentially with bits. PyTorch uses a custom autograd.Function for hard quantizers.

Parameters

NameType & defaultContract
zpositional-or-keywordTensorRequiredFloating-point input array; the quantized result has the same shape.
bitspositional-or-keywordintDefault: 2Integer number of quantization bits; creates 2**bits output levels.
lopositional-or-keywordfloatDefault: -1.0Lower endpoint of the quantization interval. Must be strictly below hi.
hipositional-or-keywordfloatDefault: 1.0Upper endpoint of the quantization interval.
betapositional-or-keywordfloatDefault: 10.0Positive finite scalar sharpness of each shifted tanh backward surrogate.
ReturnsTensor
Explicit validation & failure cases
  • When bits < 1, raises ValueError(f'bits must be >= 1, got {bits}').

These are explicit source guards. Backend operations may also raise for invalid shapes, dtypes, or unsupported devices.

Complete runnable example · binary torch

This is the complete downloadable script, tested against the pinned installed distribution. It includes initialization and assertions.

Download complete script ↓
"""Hard quantization, analytic surrogate derivatives, and an annealing schedule.

python -m pip install --pre "omnibias-binary==0.1.0a2" "omnibias-torch==0.5.0rc1"
"""
import math
import torch
from omnibias.binary import BetaAnnealScheduler
from omnibias.binary.torch.ops import (
    binarize, binarize01, heaviside, ternarize, kbit_quantize,
    riccati_tanh_derivative, riccati_sigmoid_derivative,
    surrogate_tower, surrogate_jet, curvature_corrected_slope, binarize_curvature,
)

z = torch.tensor([-0.75, -0.25, 0.0, 0.25, 0.75], dtype=torch.float64, requires_grad=True)
beta = 2.0
hard = binarize(z, beta)
torch.testing.assert_close(hard, torch.tensor([-1., -1., 1., 1., 1.], dtype=z.dtype))
# This is a specified surrogate gradient, not the classical derivative of sign.
(grad,) = torch.autograd.grad(hard.sum(), z)
torch.testing.assert_close(grad, beta * (1 - torch.tanh(beta*z).square()))
step = heaviside(z, beta)
torch.testing.assert_close(step, binarize01(z, beta))
(step_grad,) = torch.autograd.grad(step.sum(), z)
s = torch.sigmoid(beta*z)
torch.testing.assert_close(step_grad, beta*s*(1-s))
torch.testing.assert_close(ternarize(z, beta, delta=0.5), torch.tensor([-1., 0., 0., 0., 1.], dtype=z.dtype))
quantized = kbit_quantize(z, bits=2, lo=-1.0, hi=1.0, beta=beta)
assert quantized.shape == z.shape and torch.isfinite(quantized).all()
assert quantized.min() >= -1 and quantized.max() <= 1
(quantized_grad,) = torch.autograd.grad(quantized.sum(), z)
assert torch.isfinite(quantized_grad).all() and quantized_grad.abs().sum() > 0

t = torch.tanh(z)
torch.testing.assert_close(riccati_tanh_derivative(t, 1), 1-t.square())
s = torch.sigmoid(z)
torch.testing.assert_close(riccati_sigmoid_derivative(s, 1), s*(1-s))
tower = surrogate_tower(z, beta, order=3)
jet = surrogate_jet(z, beta, order=3)
assert tower.shape == jet.shape == (4, 5)
for order in range(4):
    torch.testing.assert_close(jet[order]*math.factorial(order), tower[order])
torch.testing.assert_close(curvature_corrected_slope(z, beta, window=0.0), tower[1])
(curvature_grad,) = torch.autograd.grad(binarize_curvature(z, beta).sum(), z)
torch.testing.assert_close(curvature_grad, curvature_corrected_slope(z, beta))

schedule = BetaAnnealScheduler(1.0, 4.0, num_steps=2, schedule="linear")
assert [schedule.step() for _ in range(3)] == [1.0, 2.5, 4.0]
assert schedule.current_step == 3 and schedule.value(100) == 4.0
schedule.reset()
assert schedule.current_step == 0 and schedule.step() == 1.0
print("quantizers: (5,); surrogate tower/jet: (4, 5); gradients and schedule passed")
python binary_torch.py
Executed regression · shape preservation

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗
def test_shape_preservation() -> None:
    z = torch.as_tensor(Z_2D, dtype=torch.float64)
    assert tq.binarize(z).shape == Z_2D.shape
    assert tq.ternarize(z).shape == Z_2D.shape
    assert tq.kbit_quantize(z, bits=2).shape == Z_2D.shape
python -m pytest "packages/omnibias-binary/tests/test_quantize.py::test_shape_preservation" -q
Read the implementation
def kbit_quantize(
    z: Tensor,
    bits: int = 2,
    lo: float = -1.0,
    hi: float = 1.0,
    beta: float = 10.0,
) -> Tensor:
    """Uniform k-bit quantize on ``[lo, hi]``; tanh-step surrogate backward."""
    if bits < 1:
        raise ValueError(f"bits must be >= 1, got {bits}")
    return _KBitQuantizeFn.apply(z, bits, lo, hi, beta)
Other public import paths
  • omnibias.binary.torch.ops.kbit_quantize
  • omnibias.binary.torch.ops.quantize.kbit_quantize
FUNCTION

riccati_sigmoid_derivative​

omnibias.binary.torch.ops.riccati_sigmoid_derivativeSource ↗
riccati_sigmoid_derivative(s: Tensor, order: int=1) -> Tensor

Evaluate P_order(s) = sigmoid^(order)(z) as a polynomial in s = sigmoid(z).

Evaluate the order-th derivative of sigmoid with respect to z from the activation value s=sigmoid(z). This avoids recomputing sigmoid when it is already available. Order zero returns s and order one returns s*(1-s); successive derivative polynomials come from the shared core recurrence. The result has the same shape as s. Negative orders raise ValueError. The formula is analytic, while floating-point evaluation still incurs rounding and cancellation, especially at large order or saturated inputs. PyTorch uses a custom autograd.Function for hard quantizers.

Parameters

NameType & defaultContract
spositional-or-keywordTensorRequiredSigmoid activation values, normally in [0, 1]; do not pass raw preactivation values here.
orderpositional-or-keywordintDefault: 1Nonnegative integer derivative order with respect to the original preactivation.
ReturnsTensor
Explicit validation & failure cases
  • When order < 0, raises ValueError(f'order must be >= 0, got {order}').

These are explicit source guards. Backend operations may also raise for invalid shapes, dtypes, or unsupported devices.

Complete runnable example · binary torch

This is the complete downloadable script, tested against the pinned installed distribution. It includes initialization and assertions.

Download complete script ↓
"""Hard quantization, analytic surrogate derivatives, and an annealing schedule.

python -m pip install --pre "omnibias-binary==0.1.0a2" "omnibias-torch==0.5.0rc1"
"""
import math
import torch
from omnibias.binary import BetaAnnealScheduler
from omnibias.binary.torch.ops import (
    binarize, binarize01, heaviside, ternarize, kbit_quantize,
    riccati_tanh_derivative, riccati_sigmoid_derivative,
    surrogate_tower, surrogate_jet, curvature_corrected_slope, binarize_curvature,
)

z = torch.tensor([-0.75, -0.25, 0.0, 0.25, 0.75], dtype=torch.float64, requires_grad=True)
beta = 2.0
hard = binarize(z, beta)
torch.testing.assert_close(hard, torch.tensor([-1., -1., 1., 1., 1.], dtype=z.dtype))
# This is a specified surrogate gradient, not the classical derivative of sign.
(grad,) = torch.autograd.grad(hard.sum(), z)
torch.testing.assert_close(grad, beta * (1 - torch.tanh(beta*z).square()))
step = heaviside(z, beta)
torch.testing.assert_close(step, binarize01(z, beta))
(step_grad,) = torch.autograd.grad(step.sum(), z)
s = torch.sigmoid(beta*z)
torch.testing.assert_close(step_grad, beta*s*(1-s))
torch.testing.assert_close(ternarize(z, beta, delta=0.5), torch.tensor([-1., 0., 0., 0., 1.], dtype=z.dtype))
quantized = kbit_quantize(z, bits=2, lo=-1.0, hi=1.0, beta=beta)
assert quantized.shape == z.shape and torch.isfinite(quantized).all()
assert quantized.min() >= -1 and quantized.max() <= 1
(quantized_grad,) = torch.autograd.grad(quantized.sum(), z)
assert torch.isfinite(quantized_grad).all() and quantized_grad.abs().sum() > 0

t = torch.tanh(z)
torch.testing.assert_close(riccati_tanh_derivative(t, 1), 1-t.square())
s = torch.sigmoid(z)
torch.testing.assert_close(riccati_sigmoid_derivative(s, 1), s*(1-s))
tower = surrogate_tower(z, beta, order=3)
jet = surrogate_jet(z, beta, order=3)
assert tower.shape == jet.shape == (4, 5)
for order in range(4):
    torch.testing.assert_close(jet[order]*math.factorial(order), tower[order])
torch.testing.assert_close(curvature_corrected_slope(z, beta, window=0.0), tower[1])
(curvature_grad,) = torch.autograd.grad(binarize_curvature(z, beta).sum(), z)
torch.testing.assert_close(curvature_grad, curvature_corrected_slope(z, beta))

schedule = BetaAnnealScheduler(1.0, 4.0, num_steps=2, schedule="linear")
assert [schedule.step() for _ in range(3)] == [1.0, 2.5, 4.0]
assert schedule.current_step == 3 and schedule.value(100) == 4.0
schedule.reset()
assert schedule.current_step == 0 and schedule.step() == 1.0
print("quantizers: (5,); surrogate tower/jet: (4, 5); gradients and schedule passed")
python binary_torch.py
Executed regression · backward sigmoid formula

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗
def test_backward_sigmoid_formula() -> None:
    z = torch.tensor(Z, dtype=torch.float64, requires_grad=True)
    tq.binarize01(z, beta=BETA).sum().backward()
    s = 1.0 / (1.0 + np.exp(-BETA * Z))
    assert np.allclose(_np(z.grad), BETA * s * (1.0 - s), rtol=1e-12, atol=1e-14)
python -m pytest "packages/omnibias-binary/tests/test_quantize01.py::test_backward_sigmoid_formula" -q
Read the implementation
def riccati_sigmoid_derivative(s: Tensor, order: int = 1) -> Tensor:
    """Evaluate ``P_order(s) = sigmoid^(order)(z)`` as a polynomial in ``s = sigmoid(z)``."""
    if order < 0:
        raise ValueError(f"order must be >= 0, got {order}")
    if order == 0:
        return s
    coeffs = sigmoid_polynomial_coeffs(order)
    return _horner(coeffs, s)
Related APIsomnibias.core.sigmoid_polynomial_coeffs↗
Other public import paths
  • omnibias.binary.torch.ops.riccati_sigmoid_derivative
  • omnibias.binary.torch.ops.quantize.riccati_sigmoid_derivative
FUNCTION

riccati_tanh_derivative​

omnibias.binary.torch.ops.riccati_tanh_derivativeSource ↗
riccati_tanh_derivative(t: Tensor, order: int=1) -> Tensor

Evaluate T_order(t) = tanh^(order)(z) as a polynomial in t = tanh(z).

Evaluate the order-th derivative of tanh with respect to its original argument, using an already-computed activation value t=tanh(z). Pass t, not the preactivation z: the polynomial is expressed in activation coordinates. Order zero returns t, order one returns 1-t**2, and higher orders use the shared recurrence coefficients. The output preserves t’s shape and follows backend tensor arithmetic. Negative orders raise ValueError. Mathematical polynomial identities do not guarantee numerical stability at arbitrarily high order or near cancellation; validate the selected dtype and range. PyTorch uses a custom autograd.Function for hard quantizers.

Parameters

NameType & defaultContract
tpositional-or-keywordTensorRequiredTanh activation values, normally in [-1, 1]. Values outside that range still evaluate a polynomial but no longer represent real tanh preactivations.
orderpositional-or-keywordintDefault: 1Nonnegative integer derivative order with respect to z, not a derivative with respect to t.
ReturnsTensor
Explicit validation & failure cases
  • When order < 0, raises ValueError(f'order must be >= 0, got {order}').

These are explicit source guards. Backend operations may also raise for invalid shapes, dtypes, or unsupported devices.

Complete runnable example · binary torch

This is the complete downloadable script, tested against the pinned installed distribution. It includes initialization and assertions.

Download complete script ↓
"""Hard quantization, analytic surrogate derivatives, and an annealing schedule.

python -m pip install --pre "omnibias-binary==0.1.0a2" "omnibias-torch==0.5.0rc1"
"""
import math
import torch
from omnibias.binary import BetaAnnealScheduler
from omnibias.binary.torch.ops import (
    binarize, binarize01, heaviside, ternarize, kbit_quantize,
    riccati_tanh_derivative, riccati_sigmoid_derivative,
    surrogate_tower, surrogate_jet, curvature_corrected_slope, binarize_curvature,
)

z = torch.tensor([-0.75, -0.25, 0.0, 0.25, 0.75], dtype=torch.float64, requires_grad=True)
beta = 2.0
hard = binarize(z, beta)
torch.testing.assert_close(hard, torch.tensor([-1., -1., 1., 1., 1.], dtype=z.dtype))
# This is a specified surrogate gradient, not the classical derivative of sign.
(grad,) = torch.autograd.grad(hard.sum(), z)
torch.testing.assert_close(grad, beta * (1 - torch.tanh(beta*z).square()))
step = heaviside(z, beta)
torch.testing.assert_close(step, binarize01(z, beta))
(step_grad,) = torch.autograd.grad(step.sum(), z)
s = torch.sigmoid(beta*z)
torch.testing.assert_close(step_grad, beta*s*(1-s))
torch.testing.assert_close(ternarize(z, beta, delta=0.5), torch.tensor([-1., 0., 0., 0., 1.], dtype=z.dtype))
quantized = kbit_quantize(z, bits=2, lo=-1.0, hi=1.0, beta=beta)
assert quantized.shape == z.shape and torch.isfinite(quantized).all()
assert quantized.min() >= -1 and quantized.max() <= 1
(quantized_grad,) = torch.autograd.grad(quantized.sum(), z)
assert torch.isfinite(quantized_grad).all() and quantized_grad.abs().sum() > 0

t = torch.tanh(z)
torch.testing.assert_close(riccati_tanh_derivative(t, 1), 1-t.square())
s = torch.sigmoid(z)
torch.testing.assert_close(riccati_sigmoid_derivative(s, 1), s*(1-s))
tower = surrogate_tower(z, beta, order=3)
jet = surrogate_jet(z, beta, order=3)
assert tower.shape == jet.shape == (4, 5)
for order in range(4):
    torch.testing.assert_close(jet[order]*math.factorial(order), tower[order])
torch.testing.assert_close(curvature_corrected_slope(z, beta, window=0.0), tower[1])
(curvature_grad,) = torch.autograd.grad(binarize_curvature(z, beta).sum(), z)
torch.testing.assert_close(curvature_grad, curvature_corrected_slope(z, beta))

schedule = BetaAnnealScheduler(1.0, 4.0, num_steps=2, schedule="linear")
assert [schedule.step() for _ in range(3)] == [1.0, 2.5, 4.0]
assert schedule.current_step == 3 and schedule.value(100) == 4.0
schedule.reset()
assert schedule.current_step == 0 and schedule.step() == 1.0
print("quantizers: (5,); surrogate tower/jet: (4, 5); gradients and schedule passed")
python binary_torch.py
Executed regression · surrogate jet parity

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗
def test_surrogate_jet_parity() -> None:
    tj = _np(tq.surrogate_jet(torch.as_tensor(Z), BETA, order=4))
    jj = _np(jq.surrogate_jet(jnp.asarray(Z), BETA, order=4))
    assert np.allclose(tj, jj, rtol=1e-12, atol=1e-13)
python -m pytest "packages/omnibias-binary/tests/test_surrogate.py::test_surrogate_jet_parity" -q
Read the implementation
def riccati_tanh_derivative(t: Tensor, order: int = 1) -> Tensor:
    """Evaluate ``T_order(t) = tanh^(order)(z)`` as a polynomial in ``t = tanh(z)``."""
    if order < 0:
        raise ValueError(f"order must be >= 0, got {order}")
    if order == 0:
        return t
    coeffs = tanh_polynomial_coeffs(order)
    return _horner(coeffs, t)
Related APIsomnibias.core.tanh_polynomial_coeffs↗
Other public import paths
  • omnibias.binary.torch.ops.riccati_tanh_derivative
  • omnibias.binary.torch.ops.quantize.riccati_tanh_derivative
FUNCTION

ternarize​

omnibias.binary.torch.ops.ternarizeSource ↗
ternarize(z: Tensor, beta: float=10.0, delta: float=0.5) -> Tensor

Hard ternary {-1, 0, +1}; smooth tanh dead-zone surrogate backward.

Use a hard three-level forward decision with a dead zone around zero. Values strictly above delta map to +1, values strictly below -delta map to -1, and the closed interval between them maps to 0. The backward rule averages the slopes of two shifted tanh surrogates. This is not the classical derivative of the piecewise-constant forward map. Use a nonnegative finite delta and positive finite beta; these domain assumptions are not fully validated by this wrapper. The delta threshold is fixed by the custom backward rule. PyTorch uses a custom autograd.Function for hard quantizers.

Parameters

NameType & defaultContract
zpositional-or-keywordTensorRequiredFloating-point preactivation array; shape is preserved.
betapositional-or-keywordfloatDefault: 10.0Positive finite scalar sharpness of the two backward transition regions.
deltapositional-or-keywordfloatDefault: 0.5Nonnegative half-width of the zero-output dead zone; values exactly at either boundary map to zero.
ReturnsTensor
Complete runnable example · binary torch

This is the complete downloadable script, tested against the pinned installed distribution. It includes initialization and assertions.

Download complete script ↓
"""Hard quantization, analytic surrogate derivatives, and an annealing schedule.

python -m pip install --pre "omnibias-binary==0.1.0a2" "omnibias-torch==0.5.0rc1"
"""
import math
import torch
from omnibias.binary import BetaAnnealScheduler
from omnibias.binary.torch.ops import (
    binarize, binarize01, heaviside, ternarize, kbit_quantize,
    riccati_tanh_derivative, riccati_sigmoid_derivative,
    surrogate_tower, surrogate_jet, curvature_corrected_slope, binarize_curvature,
)

z = torch.tensor([-0.75, -0.25, 0.0, 0.25, 0.75], dtype=torch.float64, requires_grad=True)
beta = 2.0
hard = binarize(z, beta)
torch.testing.assert_close(hard, torch.tensor([-1., -1., 1., 1., 1.], dtype=z.dtype))
# This is a specified surrogate gradient, not the classical derivative of sign.
(grad,) = torch.autograd.grad(hard.sum(), z)
torch.testing.assert_close(grad, beta * (1 - torch.tanh(beta*z).square()))
step = heaviside(z, beta)
torch.testing.assert_close(step, binarize01(z, beta))
(step_grad,) = torch.autograd.grad(step.sum(), z)
s = torch.sigmoid(beta*z)
torch.testing.assert_close(step_grad, beta*s*(1-s))
torch.testing.assert_close(ternarize(z, beta, delta=0.5), torch.tensor([-1., 0., 0., 0., 1.], dtype=z.dtype))
quantized = kbit_quantize(z, bits=2, lo=-1.0, hi=1.0, beta=beta)
assert quantized.shape == z.shape and torch.isfinite(quantized).all()
assert quantized.min() >= -1 and quantized.max() <= 1
(quantized_grad,) = torch.autograd.grad(quantized.sum(), z)
assert torch.isfinite(quantized_grad).all() and quantized_grad.abs().sum() > 0

t = torch.tanh(z)
torch.testing.assert_close(riccati_tanh_derivative(t, 1), 1-t.square())
s = torch.sigmoid(z)
torch.testing.assert_close(riccati_sigmoid_derivative(s, 1), s*(1-s))
tower = surrogate_tower(z, beta, order=3)
jet = surrogate_jet(z, beta, order=3)
assert tower.shape == jet.shape == (4, 5)
for order in range(4):
    torch.testing.assert_close(jet[order]*math.factorial(order), tower[order])
torch.testing.assert_close(curvature_corrected_slope(z, beta, window=0.0), tower[1])
(curvature_grad,) = torch.autograd.grad(binarize_curvature(z, beta).sum(), z)
torch.testing.assert_close(curvature_grad, curvature_corrected_slope(z, beta))

schedule = BetaAnnealScheduler(1.0, 4.0, num_steps=2, schedule="linear")
assert [schedule.step() for _ in range(3)] == [1.0, 2.5, 4.0]
assert schedule.current_step == 3 and schedule.value(100) == 4.0
schedule.reset()
assert schedule.current_step == 0 and schedule.step() == 1.0
print("quantizers: (5,); surrogate tower/jet: (4, 5); gradients and schedule passed")
python binary_torch.py
Executed regression · shape preservation

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗
def test_shape_preservation() -> None:
    z = torch.as_tensor(Z_2D, dtype=torch.float64)
    assert tq.binarize(z).shape == Z_2D.shape
    assert tq.ternarize(z).shape == Z_2D.shape
    assert tq.kbit_quantize(z, bits=2).shape == Z_2D.shape
python -m pytest "packages/omnibias-binary/tests/test_quantize.py::test_shape_preservation" -q
Read the implementation
def ternarize(z: Tensor, beta: float = 10.0, delta: float = 0.5) -> Tensor:
    """Hard ternary ``{-1, 0, +1}``; smooth ``tanh`` dead-zone surrogate backward."""
    return _TernarizeFn.apply(z, beta, delta)
Other public import paths
  • omnibias.binary.torch.ops.ternarize
  • omnibias.binary.torch.ops.quantize.ternarize
← Package overviewRunnable learning examples →