API REFERENCE/omnibias-binary

omnibias.binary.torch.ops.surrogate

0.1.0a2AlphaPython >=3.10Apache-2.0

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

import omnibias.binary.torch.ops.surrogate

Module contract & mathematical conventions

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

Bit-identical twin of omnibias.binary.jax.ops.surrogate. These go *beyond* the straight-through estimator: the hard quantizer backward uses the exact derivative tower / Taylor jet of the smooth tanh(beta z) surrogate (one tanh evaluation, any order, via the Riccati polynomials), so the backward carries the surrogate's local curvature rather than just its slope.

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.torch.ops.binarize_curvatureSource ↗

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

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

Identical hard forward to omnibias.binary.torch.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. PyTorch uses a custom autograd.Function for hard quantizers.

Parameters

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

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 · binarize curvature forward is hard sign

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_binarize_curvature_forward_is_hard_sign() -> None:
    out = _np(tq.binarize_curvature(torch.as_tensor(Z), beta=BETA))
    assert np.allclose(out, np.where(Z >= 0, 1.0, -1.0))

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

Read the implementation

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

    Identical hard forward to :func:`omnibias.binary.torch.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 _BinarizeCurvatureFn.apply(z, beta)

Other public import paths

• omnibias.binary.torch.ops.binarize_curvature
• omnibias.binary.torch.ops.surrogate.binarize_curvature

FUNCTION

curvature_corrected_slope

omnibias.binary.torch.ops.curvature_corrected_slopeSource ↗

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

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):

a better effective gradient through the hard step than the point slope alone. 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. PyTorch uses a custom autograd.Function for hard quantizers.

Parameters

Name	Type & default	Contract
zpositional-or-keyword	TensorRequired	Floating-point sample positions.
betapositional-or-keyword	float | TensorRequired	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.

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 · 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: Tensor, beta: float | Tensor, *, window: float | None = None
) -> Tensor:
    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``):

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

    a better effective gradient through the hard step than the point slope alone.
    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.torch.ops.surrogate_tower↗

Other public import paths

• omnibias.binary.torch.ops.curvature_corrected_slope
• omnibias.binary.torch.ops.surrogate.curvature_corrected_slope

FUNCTION

surrogate_jet

omnibias.binary.torch.ops.surrogate_jetSource ↗

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

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)). jet[1] is the surrogate slope (the standard backward), jet[2] its curvature: 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. PyTorch uses a custom autograd.Function for hard quantizers.

Parameters

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

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 surrogate_jet(z: Tensor, beta: float | Tensor, order: int) -> Tensor:
    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))``. ``jet[1]`` is
    the surrogate slope (the standard backward), ``jet[2]`` its curvature: this is
    the jet-STE backward signal.
    """
    if order < 0:
        raise ValueError(f"order must be >= 0, got {order}")
    z = torch.as_tensor(z)
    t = torch.tanh(beta * z)
    sigma_rows = [t] + [riccati_tanh_derivative(t, order=k) for k in range(1, order + 1)]
    sigma_tower = torch.stack(sigma_rows, dim=0)
    u_rows = [beta * z, torch.ones_like(z) * beta]
    u_rows += [torch.zeros_like(z) for _ in range(order - 1)]
    u_jet = torch.stack(u_rows[: order + 1], dim=0)
    return compose_jet(u_jet, sigma_tower)

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

Other public import paths

• omnibias.binary.torch.ops.surrogate_jet
• omnibias.binary.torch.ops.surrogate.surrogate_jet

FUNCTION

surrogate_tower

omnibias.binary.torch.ops.surrogate_towerSource ↗

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

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. PyTorch uses a custom autograd.Function for hard quantizers.

Parameters

Name	Type & default	Contract
zpositional-or-keyword	TensorRequired	Floating-point positions at which to evaluate the smooth surrogate.
betapositional-or-keyword	float | TensorRequired	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.

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 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: Tensor, beta: float | Tensor, order: int) -> Tensor:
    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 = torch.as_tensor(z)
    t = torch.tanh(beta * z)
    rows = [t]
    for k in range(1, order + 1):
        rows.append((beta**k) * riccati_tanh_derivative(t, order=k))
    return torch.stack(rows, dim=0)

Related APIsomnibias.binary.torch.ops.riccati_tanh_derivative↗

Other public import paths

• omnibias.binary.torch.ops.surrogate_tower
• omnibias.binary.torch.ops.surrogate.surrogate_tower

← Package overviewRunnable learning examples →