omnibias.discrete.torch.relaxation
Shared discrete-problem contracts, annealed relaxations, hard decoders, and optimality-gap evidence.
import omnibias.discrete.torch.relaxationModule contract & mathematical conventions
Differentiable annealed relaxation core (torch) via the temperature-collapse penalty.
Bit-identical twin of omnibias.discrete.jax.relaxation (float64); see that module
for the full math. Differentiable through autograd so a model that predicts the
problem coefficients can be trained *through* the relaxation.
Terminology: the beta -> inf hardening of sigmoid here is the feasibility /
temperature sense of "collapse" (a soft indicator becoming a 0/1 step), distinct from
the **founding bias collapse** (the multi-bias delta -> 0 limit to sigma^(K-1),
a derivative; see omnibias.torch.unit).
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.
anneal_descent
omnibias.discrete.torch.anneal_descentSource ↗anneal_descent(grad_x_fn: Callable[[Tensor], Tensor], scale: Any, n: int, schedule: AnnealSchedule | None=None) -> TensorAnneal x = sigmoid(beta theta) to a vertex by descending grad_x_fn.
See omnibias.discrete.jax.relaxation.anneal_descent for the full description;
this is the mathematically equivalent torch twin.
Backend scope: this identifies the corresponding mathematical construction. It does not guarantee bitwise-equal floating-point results across devices, compilers, dtypes, or backend versions. Validate the selected path with an appropriate absolute and relative tolerance.
Parameters
| Name | Type & default | Contract |
|---|---|---|
grad_x_fnpositional-or-keyword | Callable[[Tensor], Tensor]Required | See the operation contract and implementation below. |
scalepositional-or-keyword | AnyRequired | See the operation contract and implementation below. |
npositional-or-keyword | intRequired | See the operation contract and implementation below. |
schedulepositional-or-keyword | AnnealSchedule | NoneDefault: None | See the operation contract and implementation below. |
TensorExecuted regression · descent is differentiable in torch
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.
def test_descent_is_differentiable_in_torch() -> None:
torch = pytest.importorskip("torch")
from omnibias.discrete.torch import anneal_descent as ta
a, b, scale = _quadratic(4, n=5)
at = torch.tensor(a, dtype=torch.float64, requires_grad=True)
bt = torch.tensor(b, dtype=torch.float64, requires_grad=True)
out = ta(lambda x: 2.0 * (at @ x) + bt, scale, 5, AnnealSchedule.fast())
out.sum().backward()
assert at.grad is not None and torch.all(torch.isfinite(at.grad))
assert bt.grad is not None and torch.all(torch.isfinite(bt.grad))python -m pytest "packages/omnibias-discrete/tests/test_relaxation.py::test_descent_is_differentiable_in_torch" -qRegression example · torch jax bit identical
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.
def test_torch_jax_bit_identical() -> None:
jnp = pytest.importorskip("jax.numpy")
torch = pytest.importorskip("torch")
from omnibias.discrete.jax import anneal_descent as ja
from omnibias.discrete.torch import anneal_descent as ta
a, b, scale = _quadratic(0) # well-determined (strong linear term) -> commits cleanly
n = a.shape[0]
aj, bj = jnp.asarray(a), jnp.asarray(b)
xj = np.asarray(ja(lambda x: 2.0 * (aj @ x) + bj, scale, n))
at, bt = torch.as_tensor(a), torch.as_tensor(b)
xt = ta(lambda x: 2.0 * (at @ x) + bt, scale, n).detach().numpy()
assert np.max(np.abs(xj - xt)) < 1e-9python -m pytest packages/omnibias-discrete/tests/test_relaxation.py -k test_torch_jax_bit_identical -qRead the implementation
def anneal_descent(
grad_x_fn: Callable[[Tensor], Tensor],
scale: Any,
n: int,
schedule: AnnealSchedule | None = None,
) -> Tensor:
r"""Anneal ``x = sigmoid(beta theta)`` to a vertex by descending ``grad_x_fn``.
See :func:`omnibias.discrete.jax.relaxation.anneal_descent` for the full description;
this is the bit-identical torch twin.
"""
sched = schedule or AnnealSchedule()
theta = torch.as_tensor(initial_theta(n), dtype=torch.float64)
# Detached so the step size never perturbs the gradient w.r.t. the coefficients.
scale_c = torch.as_tensor(scale, dtype=torch.float64).detach()
betas = sched.betas()
for beta in betas:
eta = sched.step_safety / (beta * 0.25 * scale_c + 1e-30)
for _ in range(sched.steps):
x = torch.sigmoid(beta * theta)
grad_x = grad_x_fn(x)
grad_theta = grad_x * (beta * x * (1.0 - x))
theta = theta - eta * grad_theta
return torch.sigmoid(betas[-1] * theta)Other public import paths
omnibias.discrete.torch.anneal_descentomnibias.discrete.torch.relaxation.anneal_descent