omnibias.binary
Binary and multilevel quantization with explicit hard-forward and surrogate-gradient contracts.
import omnibias.binaryUse binary to investigate quantized values while training with the package's documented gradient rule. A hard quantizer and its smooth surrogate are different mathematical functions. Read which value is used in the forward pass and which derivative is used during differentiation.
The tanh-based temperature parameter changes the width and magnitude of the surrogate derivative. Hardening can concentrate learning near a threshold and suppress gradients elsewhere. It does not make the discontinuous quantizer classically differentiable at its jump.
Validate the quantization levels, zero convention, tensor shape, and parameter-gradient behavior for the selected Torch or JAX path. The package is alpha; do not infer deployment precision, hardware acceleration, or end-to-end training speed from the presence of a k-bit quantization primitive.
Module contract & mathematical conventions
omnibias-binary: closed-form quantization gradients.
Hard forward quantizers (binary / ternary / k-bit uniform) paired with
backward passes that use the exact derivative of a smooth tanh(beta * z)
surrogate via the Riccati identity tanh'(z) = 1 - tanh(z)^2. Coefficients
come from omnibias.core.polynomials.tanh_polynomial_coeffs.
Backend ops live under omnibias.binary.torch and omnibias.binary.jax.
important:
**Bit-parity with the PyTorch twin requires 64-bit JAX** -- jax.config.update("jax_enable_x64", True) before the first JAX array is created (or JAX_ENABLE_X64=1). JAX otherwise truncates to float32 while PyTorch uses float64, so the twins stay internally consistent but agree only to float32 tolerance. Where a value feeds a threshold, a rounding step or an argmax, that is enough to change the decision rather than just the last digits. See omnibias.jax.precision.
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.
Imports & aliases
These names share a definition with another public module. Follow a name for its full contract, parameters, and examples.
| Import from this module | Canonical reference |
|---|---|
BetaAnnealScheduler | omnibias.binary.BetaAnnealScheduler |