omnibias.struct.jax.dtw
Differentiable dynamic programs for sequences and parses, structured distributions, and finite-temperature bounds.
import omnibias.struct.jax.dtwModule contract & mathematical conventions
Differentiable soft Dynamic Time Warping (jax).
Bit-identical twin of omnibias.struct.torch.dtw (float64; needs jax_enable_x64).
Soft-DTW replaces the hard min of the DTW recursion with a **softmin** -lse_beta(-.)
built from the shared omnibias.struct.jax._logsumexp.logsumexp_beta, unrolled so
jax.grad flows through it. As beta -> inf (the temperature axis) it anneals to hard
DTW from below; the delta -> 0 tower differentiates it exactly
(soft_dtw_marginals is the closed-form soft-alignment matrix, pinned equal to
jax.grad). Do not conflate the two axes.
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.
soft_dtw
omnibias.struct.jax.soft_dtwSource ↗soft_dtw(cost: Array, beta: float=1.0) -> ArraySoft-DTW value -beta^-1 log sum_paths exp(-beta cost(path)) of an (n, m) cost matrix.
The recursive softmin of the DTW lattice; equals the flat softmin over all monotonic
paths and -> hard DTW as beta -> inf. Differentiable in cost.
Parameters
| Name | Type & default | Contract |
|---|---|---|
costpositional-or-keyword | ArrayRequired | See the operation contract and implementation below. |
betapositional-or-keyword | floatDefault: 1.0 | See the operation contract and implementation below. |
ArrayExecuted regression · soft dtw marginals equal grad jax
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_soft_dtw_marginals_equal_grad_jax() -> None:
rng = np.random.default_rng(4)
cost = _cost(rng, 3, 4)
cj = jnp.asarray(cost)
grad = jax.grad(lambda c: jdtw.soft_dtw(c, 3.0))(cj)
closed = jdtw.soft_dtw_marginals(cj, 3.0)
assert float(jnp.max(jnp.abs(grad - closed))) < 1e-9python -m pytest "packages/omnibias-struct/tests/test_dtw.py::test_soft_dtw_marginals_equal_grad_jax" -qRegression example · soft dtw torch jax 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.
def test_soft_dtw_torch_jax_parity() -> None:
rng = np.random.default_rng(6)
cost = _cost(rng, 4, 3)
ct, cj = torch.tensor(cost), jnp.asarray(cost)
for beta in (1.0, 8.0):
v_t = float(tdtw.soft_dtw(ct, beta))
v_j = float(jdtw.soft_dtw(cj, beta))
assert abs(v_t - v_j) < 1e-9
m_t = tdtw.soft_dtw_marginals(ct, beta).numpy()
m_j = np.asarray(jdtw.soft_dtw_marginals(cj, beta))
assert np.max(np.abs(m_t - m_j)) < 1e-9python -m pytest packages/omnibias-struct/tests/test_dtw.py -k test_soft_dtw_torch_jax_parity -qRead the implementation
def soft_dtw(cost: Array, beta: float = 1.0) -> Array:
r"""Soft-DTW value ``-beta^-1 log sum_paths exp(-beta cost(path))`` of an ``(n, m)`` cost matrix.
The recursive softmin of the DTW lattice; equals the flat softmin over all monotonic
paths and ``-> hard DTW`` as ``beta -> inf``. Differentiable in ``cost``.
"""
return _dtw_forward(cost, beta)[-1][-1]Other public import paths
omnibias.struct.jax.soft_dtwomnibias.struct.jax.dtw.soft_dtw
soft_dtw_batched
omnibias.struct.jax.soft_dtw_batchedSource ↗soft_dtw_batched(cost: Array, beta: float=1.0) -> ArrayBatched soft_dtw -> (B,) for cost (B, n, m) (via jax.vmap).
Parameters
| Name | Type & default | Contract |
|---|---|---|
costpositional-or-keyword | ArrayRequired | See the operation contract and implementation below. |
betapositional-or-keyword | floatDefault: 1.0 | See the operation contract and implementation below. |
ArrayExecuted regression · soft dtw batched matches loop
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_soft_dtw_batched_matches_loop() -> None:
rng = np.random.default_rng(7)
costs = np.stack([_cost(rng, 3, 4) for _ in range(5)])
ct, cj = torch.tensor(costs), jnp.asarray(costs)
beta = 4.0
bt = tdtw.soft_dtw_batched(ct, beta).numpy()
bj = np.asarray(jdtw.soft_dtw_batched(cj, beta))
loop = np.array([float(tdtw.soft_dtw(torch.tensor(c), beta)) for c in costs])
assert np.max(np.abs(bt - loop)) < 1e-9
assert np.max(np.abs(bj - loop)) < 1e-9python -m pytest "packages/omnibias-struct/tests/test_dtw.py::test_soft_dtw_batched_matches_loop" -qRead the implementation
def soft_dtw_batched(cost: Array, beta: float = 1.0) -> Array:
r"""Batched :func:`soft_dtw` -> ``(B,)`` for ``cost`` ``(B, n, m)`` (via ``jax.vmap``)."""
out: Array = jax.vmap(lambda c: soft_dtw(c, beta))(cost)
return outOther public import paths
omnibias.struct.jax.soft_dtw_batchedomnibias.struct.jax.dtw.soft_dtw_batched
soft_dtw_marginals
omnibias.struct.jax.soft_dtw_marginalsSource ↗soft_dtw_marginals(cost: Array, beta: float=1.0) -> ArrayClosed-form soft-alignment matrix E[i, j] = P_beta(cell (i, j) on the warping path).
Forward-backward over the DTW lattice with the tower softmax; equals
d soft_dtw / d cost (each cell's expected usage) and concentrates on the hard DTW
path as beta -> inf. The source and sink cells always have marginal 1.
Parameters
| Name | Type & default | Contract |
|---|---|---|
costpositional-or-keyword | ArrayRequired | See the operation contract and implementation below. |
betapositional-or-keyword | floatDefault: 1.0 | See the operation contract and implementation below. |
ArrayExecuted regression · soft dtw marginals equal grad jax
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_soft_dtw_marginals_equal_grad_jax() -> None:
rng = np.random.default_rng(4)
cost = _cost(rng, 3, 4)
cj = jnp.asarray(cost)
grad = jax.grad(lambda c: jdtw.soft_dtw(c, 3.0))(cj)
closed = jdtw.soft_dtw_marginals(cj, 3.0)
assert float(jnp.max(jnp.abs(grad - closed))) < 1e-9python -m pytest "packages/omnibias-struct/tests/test_dtw.py::test_soft_dtw_marginals_equal_grad_jax" -qRegression example · soft dtw marginals equal autograd torch
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_soft_dtw_marginals_equal_autograd_torch() -> None:
rng = np.random.default_rng(3)
cost = _cost(rng, 4, 3)
ct = torch.tensor(cost, requires_grad=True)
tdtw.soft_dtw(ct, 3.0).backward()
closed = tdtw.soft_dtw_marginals(torch.tensor(cost), 3.0)
assert torch.max(torch.abs(ct.grad - closed)).item() < 1e-9python -m pytest packages/omnibias-struct/tests/test_dtw.py -k test_soft_dtw_marginals_equal_autograd_torch -qRead the implementation
def soft_dtw_marginals(cost: Array, beta: float = 1.0) -> Array:
r"""Closed-form soft-alignment matrix ``E[i, j] = P_beta(cell (i, j) on the warping path)``.
Forward-backward over the DTW lattice with the tower softmax; equals
``d soft_dtw / d cost`` (each cell's expected usage) and concentrates on the hard DTW
path as ``beta -> inf``. The source and sink cells always have marginal ``1``.
"""
n, m = cost.shape
fwd = _dtw_forward(cost, beta)
bwd = _dtw_backward(cost, beta)
value = fwd[n - 1][m - 1]
rows = [
jnp.stack([jnp.exp(-beta * (fwd[i][j] + bwd[i][j] - cost[i, j] - value)) for j in range(m)])
for i in range(n)
]
return jnp.stack(rows, axis=0)Other public import paths
omnibias.struct.jax.soft_dtw_marginalsomnibias.struct.jax.dtw.soft_dtw_marginals