omnibias.struct.jax.tropical
Differentiable dynamic programs for sequences and parses, structured distributions, and finite-temperature bounds.
import omnibias.struct.jax.tropicalModule contract & mathematical conventions
Tropical homotopy twins (jax; theory 01-08).
beta -> inf is temperature collapse, not founding delta -> 0.
Do not conflate the two.
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.
relaxed_value
omnibias.struct.jax.tropical.relaxed_valueSource ↗relaxed_value(poly: TropicalLinear, x: Array, *, beta: float) -> ArrayThe module contract describes the context for this operation. Its signature, explicit validation, return expressions, and implementation are recorded below.
Parameters
| Name | Type & default | Contract |
|---|---|---|
polypositional-or-keyword | TropicalLinearRequired | See the operation contract and implementation below. |
xpositional-or-keyword | ArrayRequired | See the operation contract and implementation below. |
betakeyword-only | floatRequired | See the operation contract and implementation below. |
ArrayExecuted regression · g4 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.
def test_g4_parity() -> None:
pytest.importorskip("jax")
import jax
import jax.numpy as jnp
import torch
from omnibias.struct.jax.tropical import relaxed_value as rv_jax
from omnibias.struct.torch.tropical import relaxed_value as rv_torch
jax.config.update("jax_enable_x64", True)
poly = _poly()
x = np.array([[0.1, -0.2], [0.3, 0.4]])
t = rv_torch(poly, torch.as_tensor(x, dtype=torch.float64), beta=2.0)
j = rv_jax(poly, jnp.asarray(x, dtype=jnp.float64), beta=2.0)
np.testing.assert_allclose(t.detach().cpu().numpy(), np.asarray(j), rtol=0.0, atol=1e-12)python -m pytest "packages/omnibias-struct/tests/test_tropical.py::test_g4_parity" -qRegression example · g2 subdivision agrees with arrangement sampler
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_g2_subdivision_agrees_with_arrangement_sampler() -> None:
poly = TropicalLinear(
np.array([0.0, 0.2, -0.1, 0.4]),
np.array([[1.0, 0.0], [0.0, 1.0], [-1.0, 0.0], [0.2, 0.8]]),
)
cells = dual_subdivision(poly, n_samples=8000, seed=0)
assert len(cells) >= 2
# Difference arrangement of the tropical terms.
n = poly.n
normals = []
offsets = []
for i in range(n):
for j in range(i + 1, n):
normals.append(poly.exponents[i] - poly.exponents[j])
offsets.append(poly.coeffs[j] - poly.coeffs[i])
arr = Arrangement(np.asarray(normals), np.asarray(offsets))
arr_cells = enumerate_cells_vertices(arr)
# Each tropical full-dim cell sits inside some arrangement cell; counts are comparable.
assert len(cells) <= len(arr_cells)
verts = newton_polytope(poly)
assert len(verts) >= 3python -m pytest packages/omnibias-struct/tests/test_tropical.py -k test_g2_subdivision_agrees_with_arrangement_sampler -qRead the implementation
def relaxed_value(poly: TropicalLinear, x: Array, *, beta: float) -> Array:
return logsumexp_beta(scores(poly, x), beta, axis=-1)relaxed_weights
omnibias.struct.jax.tropical.relaxed_weightsSource ↗relaxed_weights(poly: TropicalLinear, x: Array, *, beta: float) -> ArrayThe module contract describes the context for this operation. Its signature, explicit validation, return expressions, and implementation are recorded below.
Parameters
| Name | Type & default | Contract |
|---|---|---|
polypositional-or-keyword | TropicalLinearRequired | See the operation contract and implementation below. |
xpositional-or-keyword | ArrayRequired | See the operation contract and implementation below. |
betakeyword-only | floatRequired | See the operation contract and implementation below. |
ArrayRead the implementation
def relaxed_weights(poly: TropicalLinear, x: Array, *, beta: float) -> Array:
return softmax_beta(scores(poly, x), beta, axis=-1)scores
omnibias.struct.jax.tropical.scoresSource ↗scores(poly: TropicalLinear, x: Array) -> ArrayThe module contract describes the context for this operation. Its signature, explicit validation, return expressions, and implementation are recorded below.
Parameters
| Name | Type & default | Contract |
|---|---|---|
polypositional-or-keyword | TropicalLinearRequired | See the operation contract and implementation below. |
xpositional-or-keyword | ArrayRequired | See the operation contract and implementation below. |
ArrayExecuted regression · g4 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.
def test_g4_parity() -> None:
pytest.importorskip("jax")
import jax
import jax.numpy as jnp
import torch
from omnibias.struct.jax.tropical import relaxed_value as rv_jax
from omnibias.struct.torch.tropical import relaxed_value as rv_torch
jax.config.update("jax_enable_x64", True)
poly = _poly()
x = np.array([[0.1, -0.2], [0.3, 0.4]])
t = rv_torch(poly, torch.as_tensor(x, dtype=torch.float64), beta=2.0)
j = rv_jax(poly, jnp.asarray(x, dtype=jnp.float64), beta=2.0)
np.testing.assert_allclose(t.detach().cpu().numpy(), np.asarray(j), rtol=0.0, atol=1e-12)python -m pytest "packages/omnibias-struct/tests/test_tropical.py::test_g4_parity" -qRead the implementation
def scores(poly: TropicalLinear, x: Array) -> Array:
m = jnp.asarray(poly.exponents, dtype=x.dtype)
a = jnp.asarray(poly.coeffs, dtype=x.dtype)
xv = x if x.ndim == 2 else x[None, :]
return a + xv @ m.T