omnibias.discrete.jax.relaxation
Shared discrete-problem contracts, annealed relaxations, hard decoders, and optimality-gap evidence.
import omnibias.discrete.jax.relaxationModule contract & mathematical conventions
Differentiable annealed relaxation core (JAX) via the temperature-collapse penalty.
Bit-identical twin of omnibias.discrete.torch.relaxation (float64). Given a
caller-supplied **closed-form** energy gradient grad_x E(x), the layer parametrises a
soft assignment x = sigmoid(beta theta) in (0, 1)^n and descends
\nabla_\theta E = \bigl(\nabla_x E\bigr)\odot\bigl(\beta\, x (1 - x)\bigr),by unrolled gradient descent along a geometric beta homotopy; as beta grows the
soft assignment collapses onto a binary vertex. One jit-able, jax.grad-friendly
call, so a model that predicts the problem coefficients can be trained *through* the
relaxation. Callers supply the
gradient and a step scale; the returned soft assignment is decoded with
omnibias.discrete.decode and the gap certified with
omnibias.discrete.certify_gap.
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 docs/derivatives.md and 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.jax.anneal_descentSource ↗anneal_descent(grad_x_fn: Callable[[Array], Array], scale: Any, n: int, schedule: AnnealSchedule | None=None) -> ArrayAnneal x = sigmoid(beta theta) to a vertex by descending grad_x_fn.
Parameters
grad_x_fn: The closed-form energy gradient x -> grad_x E(x) (an (n,) array). It may close over differentiable coefficients so the whole descent is jax.grad-differentiable through them.
scale: A (detached) Lipschitz-like magnitude for the energy gradient, setting the step size eta = step_safety / (beta * 0.25 * scale). Larger is safer (slower).
n: Number of variables.
schedule: The AnnealSchedule (defaults to AnnealSchedule()).
Parameters
| Name | Type & default | Contract |
|---|---|---|
grad_x_fnpositional-or-keyword | Callable[[Array], Array]Required | The closed-form energy gradient x -> grad_x E(x) (an (n,) array). It may close over differentiable coefficients so the whole descent is jax.grad-differentiable through them. |
scalepositional-or-keyword | AnyRequired | A (detached) Lipschitz-like magnitude for the energy gradient, setting the step size eta = step_safety / (beta * 0.25 * scale). Larger is safer (slower). |
npositional-or-keyword | intRequired | Number of variables. |
schedulepositional-or-keyword | AnnealSchedule | NoneDefault: None | The AnnealSchedule (defaults to AnnealSchedule()). |
ArrayExecuted regression · output is in the unit box
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_output_is_in_the_unit_box() -> None:
jnp = pytest.importorskip("jax.numpy")
from omnibias.discrete.jax import anneal_descent as ja
a, b, scale = _quadratic(2)
aj, bj = jnp.asarray(a), jnp.asarray(b)
x = np.asarray(ja(lambda z: 2.0 * (aj @ z) + bj, scale, a.shape[0]))
assert np.all(x >= 0.0) and np.all(x <= 1.0) and np.all(np.isfinite(x))python -m pytest "packages/omnibias-discrete/tests/test_relaxation.py::test_output_is_in_the_unit_box" -qRegression example · annealing collapses toward a vertex
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_annealing_collapses_toward_a_vertex() -> None:
jnp = pytest.importorskip("jax.numpy")
from omnibias.discrete.jax import anneal_descent as ja
a, b, scale = _quadratic(0)
aj, bj = jnp.asarray(a), jnp.asarray(b)
x = np.asarray(ja(lambda z: 2.0 * (aj @ z) + bj, scale, a.shape[0], AnnealSchedule(stages=16)))
assert float(np.mean(np.minimum(x, 1.0 - x))) < 0.1python -m pytest packages/omnibias-discrete/tests/test_relaxation.py -k test_annealing_collapses_toward_a_vertex -qRead the implementation
def anneal_descent(
grad_x_fn: Callable[[Array], Array],
scale: Any,
n: int,
schedule: AnnealSchedule | None = None,
) -> Array:
r"""Anneal ``x = sigmoid(beta theta)`` to a vertex by descending ``grad_x_fn``.
Parameters
----------
grad_x_fn:
The closed-form energy gradient ``x -> grad_x E(x)`` (an ``(n,)`` array). It may
close over differentiable coefficients so the whole descent is
``jax.grad``-differentiable through them.
scale:
A (detached) Lipschitz-like magnitude for the energy gradient, setting the step
size ``eta = step_safety / (beta * 0.25 * scale)``. Larger is safer (slower).
n:
Number of variables.
schedule:
The :class:`AnnealSchedule` (defaults to ``AnnealSchedule()``).
"""
sched = schedule or AnnealSchedule()
theta = jnp.asarray(initial_theta(n), dtype=jnp.float64)
# Detached so the step size never perturbs the gradient w.r.t. the coefficients.
scale_c = jax.lax.stop_gradient(jnp.asarray(scale, dtype=jnp.float64))
def stage(theta: Array, beta: float) -> Array:
eta = sched.step_safety / (beta * 0.25 * scale_c + 1e-30)
def body(_: int, th: Array) -> Array:
x = jax.nn.sigmoid(beta * th)
grad_x = grad_x_fn(x)
grad_theta = grad_x * (beta * x * (1.0 - x))
return cast(Array, th - eta * grad_theta)
return cast(Array, jax.lax.fori_loop(0, sched.steps, body, theta))
betas = sched.betas()
for beta in betas:
theta = stage(theta, beta)
result: Array = jax.nn.sigmoid(betas[-1] * theta)
return resultOther public import paths
omnibias.discrete.jax.anneal_descentomnibias.discrete.jax.relaxation.anneal_descent