API REFERENCE/omnibias-binary

omnibias.binary.jax.ops.surrogate

0.1.0a2AlphaPython >=3.10Apache-2.0

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

import omnibias.binary.jax.ops.surrogate

Module contract & mathematical conventions

JAX smooth-surrogate towers / jets and a curvature-aware quantizer backward.

These go *beyond* the straight-through estimator: instead of treating the hard quantizer's backward as the identity, omnibias uses the exact derivative tower of the smooth tanh(beta z) surrogate (one tanh evaluation, any order, via the Riccati polynomials). The full Taylor jet of the surrogate -- built with the same omnibias.jax.jet.compose_jet machinery as Phase 1 -- carries the local *curvature* of the surrogate, not just its slope, enabling 2nd-order-corrected ("jet-STE") quantizer gradients.

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_curvature

omnibias.binary.jax.ops.binarize_curvatureSource ↗

binarize_curvature(z: Array, beta: float=10.0) -> Array

Hard sign(z) with a 2nd-order, curvature-corrected surrogate backward.

Identical hard forward to omnibias.binary.jax.ops.binarize; the backward uses curvature_corrected_slope (the windowed-average slope) instead of the point slope. beta is a fixed hyperparameter here.

Choose the curvature-corrected hard-sign training rule when an explicit window-averaged surrogate is appropriate. Forward values are identical to binarize, including +1 at zero. Reverse mode with respect to z uses curvature_corrected_slope with its default window, rather than the pointwise tanh slope. This custom rule treats beta as fixed and supplies no beta cotangent. The correction is a truncated approximation to an average gradient and is not an exact Hessian or a universal optimization improvement. The example compares the custom cotangent directly with its defining slope expression. JAX hard quantizers use custom VJPs; use reverse-mode transforms, since a custom VJP does not provide a general forward-mode JVP rule.

Parameters

Name	Type & default	Contract
zpositional-or-keyword	ArrayRequired	Floating-point preactivations of arbitrary shape.
betapositional-or-keyword	floatDefault: 10.0	Positive finite scalar hyperparameter; this custom backward does not train beta.

ReturnsArray

Complete runnable example · binary jax

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

Download complete script ↓

"""JAX hard quantization with explicitly tested surrogate cotangents.

python -m pip install --pre "omnibias-binary==0.1.0a2" "omnibias-jax==0.5.0rc1"
"""
import math
import jax
jax.config.update("jax_enable_x64", True)
import jax.numpy as jnp
from omnibias.binary.jax.ops import (
    binarize, binarize01, heaviside, ternarize, kbit_quantize,
    riccati_tanh_derivative, riccati_sigmoid_derivative,
    surrogate_tower, surrogate_jet, curvature_corrected_slope, binarize_curvature,
)

z = jnp.array([-0.75, -0.25, 0.0, 0.25, 0.75], dtype=jnp.float64)
beta = 2.0
assert jnp.array_equal(binarize(z, beta), jnp.array([-1., -1., 1., 1., 1.]))
grad = jax.grad(lambda x: binarize(x, beta).sum())(z)
assert jnp.allclose(grad, beta*(1-jnp.tanh(beta*z)**2), rtol=1e-12, atol=1e-12)
assert jnp.array_equal(heaviside(z, beta), binarize01(z, beta))
s = jax.nn.sigmoid(beta*z)
step_grad = jax.grad(lambda x: heaviside(x, beta).sum())(z)
assert jnp.allclose(step_grad, beta*s*(1-s), rtol=1e-12, atol=1e-12)
assert jnp.array_equal(ternarize(z, beta, delta=0.5), jnp.array([-1., 0., 0., 0., 1.]))
q = kbit_quantize(z, bits=2, lo=-1., hi=1., beta=beta)
assert q.shape == z.shape and jnp.isfinite(q).all() and q.min() >= -1 and q.max() <= 1
q_grad = jax.grad(lambda x: kbit_quantize(x, bits=2, beta=beta).sum())(z)
assert jnp.isfinite(q_grad).all() and jnp.abs(q_grad).sum() > 0

t = jnp.tanh(z)
assert jnp.allclose(riccati_tanh_derivative(t, 1), 1-t*t)
s = jax.nn.sigmoid(z)
assert jnp.allclose(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):
    assert jnp.allclose(jet[order]*math.factorial(order), tower[order], rtol=1e-12, atol=1e-12)
assert jnp.allclose(curvature_corrected_slope(z, beta, window=0.), tower[1])
curvature_grad = jax.grad(lambda x: binarize_curvature(x, beta).sum())(z)
assert jnp.allclose(curvature_grad, curvature_corrected_slope(z, beta), rtol=1e-12, atol=1e-12)
print("quantizers: (5,); surrogate tower/jet: (4, 5); surrogate gradients passed")

python binary_jax.py

Regression example · binarize curvature parity

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_binarize_curvature_parity() -> None:
    zt = torch.tensor(Z, dtype=torch.float64, requires_grad=True)
    tq.binarize_curvature(zt, beta=BETA).sum().backward()
    tgrad = _np(zt.grad)
    jgrad = _np(jax.grad(lambda zz: jnp.sum(jq.binarize_curvature(zz, BETA)))(jnp.asarray(Z)))
    assert np.allclose(tgrad, jgrad, rtol=1e-11, atol=1e-12)

python -m pytest packages/omnibias-binary/tests/test_surrogate.py -k test_binarize_curvature_parity -q

Read the implementation

def binarize_curvature(z: Array, beta: float = 10.0) -> Array:
    """Hard ``sign(z)`` with a 2nd-order, curvature-corrected surrogate backward.

    Identical hard forward to :func:`omnibias.binary.jax.ops.binarize`; the
    backward uses :func:`curvature_corrected_slope` (the windowed-average slope)
    instead of the point slope. ``beta`` is a fixed hyperparameter here.
    """
    return jnp.where(z >= 0, 1.0, -1.0)

Other public import paths

• omnibias.binary.jax.ops.binarize_curvature
• omnibias.binary.jax.ops.surrogate.binarize_curvature

FUNCTION

curvature_corrected_slope

omnibias.binary.jax.ops.curvature_corrected_slopeSource ↗

curvature_corrected_slope(z: Array, beta: float | Array, *, window: float | None=None) -> Array

Windowed-average surrogate slope s'(z) + (h^2/6) s'''(z) (h = window).

The point slope s'(z) is replaced by the average of s' over a symmetric window of half-width h (default h = 1/beta, the surrogate's natural smoothing scale):

Using the exact 3rd-order tower this is a better *effective* gradient through the hard step than the point slope alone (it accounts for the surrogate's finite transition width). Reduces to s'(z) as h -> 0.

Evaluate the truncated symmetric-window average s'(z) + window**2 * s'''(z)/6 for s=tanh(beta*z). Use it to inspect the alternative backward signal before choosing binarize_curvature. A zero window reduces to the ordinary surrogate slope; omitting window selects 1/beta and therefore requires nonzero beta. This is a second-order approximation to the window average, with an omitted fourth-order remainder, and is not guaranteed positive or more accurate for every window. It preserves the shape of z and does not quantize the forward value. JAX hard quantizers use custom VJPs; use reverse-mode transforms, since a custom VJP does not provide a general forward-mode JVP rule.

Parameters

Name	Type & default	Contract
zpositional-or-keyword	ArrayRequired	Floating-point sample positions.
betapositional-or-keyword	float | ArrayRequired	Positive finite scalar inverse temperature; used to choose the default window as well as the surrogate.
windowkeyword-only	float | NoneDefault: None	Optional scalar half-width. None uses 1/beta; zero removes the curvature correction.

ReturnsArray

Complete runnable example · binary jax

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

Download complete script ↓

"""JAX hard quantization with explicitly tested surrogate cotangents.

python -m pip install --pre "omnibias-binary==0.1.0a2" "omnibias-jax==0.5.0rc1"
"""
import math
import jax
jax.config.update("jax_enable_x64", True)
import jax.numpy as jnp
from omnibias.binary.jax.ops import (
    binarize, binarize01, heaviside, ternarize, kbit_quantize,
    riccati_tanh_derivative, riccati_sigmoid_derivative,
    surrogate_tower, surrogate_jet, curvature_corrected_slope, binarize_curvature,
)

z = jnp.array([-0.75, -0.25, 0.0, 0.25, 0.75], dtype=jnp.float64)
beta = 2.0
assert jnp.array_equal(binarize(z, beta), jnp.array([-1., -1., 1., 1., 1.]))
grad = jax.grad(lambda x: binarize(x, beta).sum())(z)
assert jnp.allclose(grad, beta*(1-jnp.tanh(beta*z)**2), rtol=1e-12, atol=1e-12)
assert jnp.array_equal(heaviside(z, beta), binarize01(z, beta))
s = jax.nn.sigmoid(beta*z)
step_grad = jax.grad(lambda x: heaviside(x, beta).sum())(z)
assert jnp.allclose(step_grad, beta*s*(1-s), rtol=1e-12, atol=1e-12)
assert jnp.array_equal(ternarize(z, beta, delta=0.5), jnp.array([-1., 0., 0., 0., 1.]))
q = kbit_quantize(z, bits=2, lo=-1., hi=1., beta=beta)
assert q.shape == z.shape and jnp.isfinite(q).all() and q.min() >= -1 and q.max() <= 1
q_grad = jax.grad(lambda x: kbit_quantize(x, bits=2, beta=beta).sum())(z)
assert jnp.isfinite(q_grad).all() and jnp.abs(q_grad).sum() > 0

t = jnp.tanh(z)
assert jnp.allclose(riccati_tanh_derivative(t, 1), 1-t*t)
s = jax.nn.sigmoid(z)
assert jnp.allclose(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):
    assert jnp.allclose(jet[order]*math.factorial(order), tower[order], rtol=1e-12, atol=1e-12)
assert jnp.allclose(curvature_corrected_slope(z, beta, window=0.), tower[1])
curvature_grad = jax.grad(lambda x: binarize_curvature(x, beta).sum())(z)
assert jnp.allclose(curvature_grad, curvature_corrected_slope(z, beta), rtol=1e-12, atol=1e-12)
print("quantizers: (5,); surrogate tower/jet: (4, 5); surrogate gradients passed")

python binary_jax.py

Executed regression · curvature corrected slope 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_curvature_corrected_slope_parity() -> None:
    tc = _np(tq.curvature_corrected_slope(torch.as_tensor(Z), BETA))
    jc = _np(jq.curvature_corrected_slope(jnp.asarray(Z), BETA))
    assert np.allclose(tc, jc, rtol=1e-12, atol=1e-13)

python -m pytest "packages/omnibias-binary/tests/test_surrogate.py::test_curvature_corrected_slope_parity" -q

Read the implementation

def curvature_corrected_slope(
    z: Array, beta: float | Array, *, window: float | None = None
) -> Array:
    r"""Windowed-average surrogate slope ``s'(z) + (h^2/6) s'''(z)`` (``h = window``).

    The point slope ``s'(z)`` is replaced by the average of ``s'`` over a
    symmetric window of half-width ``h`` (default ``h = 1/beta``, the surrogate's
    natural smoothing scale):

    .. math:: \frac{1}{2h}\int_{-h}^{h} s'(z + u)\,du = s'(z) + \frac{h^2}{6} s'''(z) + O(h^4).

    Using the exact 3rd-order tower this is a better *effective* gradient through
    the hard step than the point slope alone (it accounts for the surrogate's
    finite transition width). Reduces to ``s'(z)`` as ``h -> 0``.
    """
    h = (1.0 / beta) if window is None else window
    tower = surrogate_tower(z, beta, order=3)
    return tower[1] + (h * h / 6.0) * tower[3]

Related APIsomnibias.binary.jax.ops.surrogate_tower↗

Other public import paths

• omnibias.binary.jax.ops.curvature_corrected_slope
• omnibias.binary.jax.ops.surrogate.curvature_corrected_slope

FUNCTION

surrogate_jet

omnibias.binary.jax.ops.surrogate_jetSource ↗

surrogate_jet(z: Array, beta: float | Array, order: int) -> Array

Taylor jet of u -> tanh(beta u) at u = z via compose_jet.

Built by composing the exact tanh derivative tower onto the affine pre-activation jet of beta * (z + t); because the inner map is affine the result equals tower_to_jet(surrogate_tower(z, beta, order)) -- a useful cross-check exercised by the test-suite. jet[1] is the surrogate slope (the standard backward), jet[2] its curvature, and so on: this is the jet-STE backward signal.

Use factorial-normalized coefficients when composing Taylor series of the smooth surrogate. Row k is the k-th derivative of tanh(beta*z) divided by k!, and the output shape is (order+1, *z.shape). In particular, row 2 is half the second derivative, not the raw curvature. The example verifies the entire coefficient-to-tower conversion through order three. The implementation composes an affine input jet with the activation tower; all derivative-order and floating-point limitations still apply. This jet describes the smooth surrogate rather than the discontinuous hard forward operation. JAX hard quantizers use custom VJPs; use reverse-mode transforms, since a custom VJP does not provide a general forward-mode JVP rule.

Parameters

Name	Type & default	Contract
zpositional-or-keyword	ArrayRequired	Floating-point base positions for the Taylor expansion.
betapositional-or-keyword	float | ArrayRequired	Scalar inverse temperature used in tanh(beta*z).
orderpositional-or-keyword	intRequired	Nonnegative integer maximum Taylor degree.

ReturnsArray

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 jax

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

Download complete script ↓

"""JAX hard quantization with explicitly tested surrogate cotangents.

python -m pip install --pre "omnibias-binary==0.1.0a2" "omnibias-jax==0.5.0rc1"
"""
import math
import jax
jax.config.update("jax_enable_x64", True)
import jax.numpy as jnp
from omnibias.binary.jax.ops import (
    binarize, binarize01, heaviside, ternarize, kbit_quantize,
    riccati_tanh_derivative, riccati_sigmoid_derivative,
    surrogate_tower, surrogate_jet, curvature_corrected_slope, binarize_curvature,
)

z = jnp.array([-0.75, -0.25, 0.0, 0.25, 0.75], dtype=jnp.float64)
beta = 2.0
assert jnp.array_equal(binarize(z, beta), jnp.array([-1., -1., 1., 1., 1.]))
grad = jax.grad(lambda x: binarize(x, beta).sum())(z)
assert jnp.allclose(grad, beta*(1-jnp.tanh(beta*z)**2), rtol=1e-12, atol=1e-12)
assert jnp.array_equal(heaviside(z, beta), binarize01(z, beta))
s = jax.nn.sigmoid(beta*z)
step_grad = jax.grad(lambda x: heaviside(x, beta).sum())(z)
assert jnp.allclose(step_grad, beta*s*(1-s), rtol=1e-12, atol=1e-12)
assert jnp.array_equal(ternarize(z, beta, delta=0.5), jnp.array([-1., 0., 0., 0., 1.]))
q = kbit_quantize(z, bits=2, lo=-1., hi=1., beta=beta)
assert q.shape == z.shape and jnp.isfinite(q).all() and q.min() >= -1 and q.max() <= 1
q_grad = jax.grad(lambda x: kbit_quantize(x, bits=2, beta=beta).sum())(z)
assert jnp.isfinite(q_grad).all() and jnp.abs(q_grad).sum() > 0

t = jnp.tanh(z)
assert jnp.allclose(riccati_tanh_derivative(t, 1), 1-t*t)
s = jax.nn.sigmoid(z)
assert jnp.allclose(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):
    assert jnp.allclose(jet[order]*math.factorial(order), tower[order], rtol=1e-12, atol=1e-12)
assert jnp.allclose(curvature_corrected_slope(z, beta, window=0.), tower[1])
curvature_grad = jax.grad(lambda x: binarize_curvature(x, beta).sum())(z)
assert jnp.allclose(curvature_grad, curvature_corrected_slope(z, beta), rtol=1e-12, atol=1e-12)
print("quantizers: (5,); surrogate tower/jet: (4, 5); surrogate gradients passed")

python binary_jax.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 surrogate_jet(z: Array, beta: float | Array, order: int) -> Array:
    r"""Taylor jet of ``u -> tanh(beta u)`` at ``u = z`` via :func:`compose_jet`.

    Built by composing the exact ``tanh`` derivative tower onto the affine
    pre-activation jet of ``beta * (z + t)``; because the inner map is affine the
    result equals ``tower_to_jet(surrogate_tower(z, beta, order))`` -- a useful
    cross-check exercised by the test-suite. ``jet[1]`` is the surrogate slope
    (the standard backward), ``jet[2]`` its curvature, and so on: this is the
    jet-STE backward signal.
    """
    if order < 0:
        raise ValueError(f"order must be >= 0, got {order}")
    z = jnp.asarray(z)
    t = jnp.tanh(beta * z)
    # sigma^(k)(u0) for tanh at u0 = beta z (derivative w.r.t. tanh's argument).
    sigma_rows = [t] + [riccati_tanh_derivative(t, order=k) for k in range(1, order + 1)]
    sigma_tower = jnp.stack(sigma_rows, axis=0)
    # Pre-activation jet of u(t) = beta * (z + t): [beta z, beta, 0, ...].
    u_rows = [beta * z, jnp.full_like(z, 1.0) * beta]
    u_rows += [jnp.zeros_like(z) for _ in range(order - 1)]
    u_jet = jnp.stack(u_rows[: order + 1], axis=0)
    return compose_jet(u_jet, sigma_tower)

Related APIsomnibias.binary.jax.ops.riccati_tanh_derivative↗omnibias.jax.compose_jet↗

Other public import paths

• omnibias.binary.jax.ops.surrogate_jet
• omnibias.binary.jax.ops.surrogate.surrogate_jet

FUNCTION

surrogate_tower

omnibias.binary.jax.ops.surrogate_towerSource ↗

surrogate_tower(z: Array, beta: float | Array, order: int) -> Array

Derivative tower [s, s', ..., s^(order)] of s(z) = tanh(beta z).

Row k is d^k/dz^k tanh(beta z) = beta^k * T_k(tanh(beta z)) where T_k is the Riccati tanh polynomial; a single tanh evaluation suffices regardless of order. Differentiable in both z and beta.

Inspect raw spatial derivatives of the smooth training surrogate tanh(beta*z). The leading axis enumerates orders 0 through order, so an input shaped (...,) produces (order+1, ...). Row k includes the chain-rule factor beta**k and is a raw derivative, without factorial normalization. These are derivatives of the smooth surrogate, not derivatives of the hard quantizer. Tensor operations retain gradients into supported floating-point inputs. High beta and high order amplify coefficients and can worsen conditioning; test finite values and tolerances in the intended dtype. JAX hard quantizers use custom VJPs; use reverse-mode transforms, since a custom VJP does not provide a general forward-mode JVP rule.

Parameters

Name	Type & default	Contract
zpositional-or-keyword	ArrayRequired	Floating-point positions at which to evaluate the smooth surrogate.
betapositional-or-keyword	float | ArrayRequired	Scalar inverse temperature; it multiplies z before tanh.
orderpositional-or-keyword	intRequired	Nonnegative integer maximum raw derivative order; included as the final leading-axis row.

ReturnsArray

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 jax

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

Download complete script ↓

"""JAX hard quantization with explicitly tested surrogate cotangents.

python -m pip install --pre "omnibias-binary==0.1.0a2" "omnibias-jax==0.5.0rc1"
"""
import math
import jax
jax.config.update("jax_enable_x64", True)
import jax.numpy as jnp
from omnibias.binary.jax.ops import (
    binarize, binarize01, heaviside, ternarize, kbit_quantize,
    riccati_tanh_derivative, riccati_sigmoid_derivative,
    surrogate_tower, surrogate_jet, curvature_corrected_slope, binarize_curvature,
)

z = jnp.array([-0.75, -0.25, 0.0, 0.25, 0.75], dtype=jnp.float64)
beta = 2.0
assert jnp.array_equal(binarize(z, beta), jnp.array([-1., -1., 1., 1., 1.]))
grad = jax.grad(lambda x: binarize(x, beta).sum())(z)
assert jnp.allclose(grad, beta*(1-jnp.tanh(beta*z)**2), rtol=1e-12, atol=1e-12)
assert jnp.array_equal(heaviside(z, beta), binarize01(z, beta))
s = jax.nn.sigmoid(beta*z)
step_grad = jax.grad(lambda x: heaviside(x, beta).sum())(z)
assert jnp.allclose(step_grad, beta*s*(1-s), rtol=1e-12, atol=1e-12)
assert jnp.array_equal(ternarize(z, beta, delta=0.5), jnp.array([-1., 0., 0., 0., 1.]))
q = kbit_quantize(z, bits=2, lo=-1., hi=1., beta=beta)
assert q.shape == z.shape and jnp.isfinite(q).all() and q.min() >= -1 and q.max() <= 1
q_grad = jax.grad(lambda x: kbit_quantize(x, bits=2, beta=beta).sum())(z)
assert jnp.isfinite(q_grad).all() and jnp.abs(q_grad).sum() > 0

t = jnp.tanh(z)
assert jnp.allclose(riccati_tanh_derivative(t, 1), 1-t*t)
s = jax.nn.sigmoid(z)
assert jnp.allclose(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):
    assert jnp.allclose(jet[order]*math.factorial(order), tower[order], rtol=1e-12, atol=1e-12)
assert jnp.allclose(curvature_corrected_slope(z, beta, window=0.), tower[1])
curvature_grad = jax.grad(lambda x: binarize_curvature(x, beta).sum())(z)
assert jnp.allclose(curvature_grad, curvature_corrected_slope(z, beta), rtol=1e-12, atol=1e-12)
print("quantizers: (5,); surrogate tower/jet: (4, 5); surrogate gradients passed")

python binary_jax.py

Executed regression · surrogate tower 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_tower_parity() -> None:
    tt = _np(tq.surrogate_tower(torch.as_tensor(Z), BETA, order=4))
    jt = _np(jq.surrogate_tower(jnp.asarray(Z), BETA, order=4))
    assert np.allclose(tt, jt, rtol=1e-12, atol=1e-13)

python -m pytest "packages/omnibias-binary/tests/test_surrogate.py::test_surrogate_tower_parity" -q

Read the implementation

def surrogate_tower(z: Array, beta: float | Array, order: int) -> Array:
    r"""Derivative tower ``[s, s', ..., s^(order)]`` of ``s(z) = tanh(beta z)``.

    Row ``k`` is ``d^k/dz^k tanh(beta z) = beta^k * T_k(tanh(beta z))`` where
    ``T_k`` is the Riccati tanh polynomial; a single ``tanh`` evaluation suffices
    regardless of order. Differentiable in both ``z`` and ``beta``.
    """
    if order < 0:
        raise ValueError(f"order must be >= 0, got {order}")
    z = jnp.asarray(z)
    t = jnp.tanh(beta * z)
    rows = [t]
    for k in range(1, order + 1):
        rows.append((beta**k) * riccati_tanh_derivative(t, order=k))
    return jnp.stack(rows, axis=0)

Related APIsomnibias.binary.jax.ops.riccati_tanh_derivative↗

Other public import paths

• omnibias.binary.jax.ops.surrogate_tower
• omnibias.binary.jax.ops.surrogate.surrogate_tower

← Package overviewRunnable learning examples →