omnibias.struct.torch.dtw
Differentiable dynamic programs for sequences and parses, structured distributions, and finite-temperature bounds.
import omnibias.struct.torch.dtwModule contract & mathematical conventions
Differentiable soft Dynamic Time Warping (torch).
Bit-identical twin of omnibias.struct.jax.dtw (float64). Soft-DTW replaces the hard
min of the DTW recursion with a **softmin** -lse_beta(-.) built from the shared
omnibias.struct.torch._logsumexp.logsumexp_beta primitive, unrolled so autograd
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 d value / d cost, pinned equal to autograd). Use
float64. 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.torch.soft_dtwSource ↗soft_dtw(cost: Tensor, beta: float=1.0) -> TensorSoft-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 | TensorRequired | See the operation contract and implementation below. |
betapositional-or-keyword | floatDefault: 1.0 | See the operation contract and implementation below. |
TensorExecuted regression · soft dtw marginals equal autograd 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_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::test_soft_dtw_marginals_equal_autograd_torch" -qRegression example · recursive soft dtw equals flat softmin over paths
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_recursive_soft_dtw_equals_flat_softmin_over_paths() -> None:
rng = np.random.default_rng(1)
cost = _cost(rng, 4, 3)
ct = torch.tensor(cost)
for beta in (0.5, 1.0, 4.0, 16.0):
recursive = float(tdtw.soft_dtw(ct, beta))
brute = brute_force_soft_dtw(cost, beta)
assert abs(recursive - brute) < 1e-9python -m pytest packages/omnibias-struct/tests/test_dtw.py -k test_recursive_soft_dtw_equals_flat_softmin_over_paths -qRead the implementation
def soft_dtw(cost: Tensor, beta: float = 1.0) -> Tensor:
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.torch.soft_dtwomnibias.struct.torch.dtw.soft_dtw
soft_dtw_batched
omnibias.struct.torch.soft_dtw_batchedSource ↗soft_dtw_batched(cost: Tensor, beta: float=1.0) -> TensorBatched soft_dtw -> (B,) for cost (B, n, m) (via torch.func.vmap).
Parameters
| Name | Type & default | Contract |
|---|---|---|
costpositional-or-keyword | TensorRequired | See the operation contract and implementation below. |
betapositional-or-keyword | floatDefault: 1.0 | See the operation contract and implementation below. |
TensorExecuted 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: Tensor, beta: float = 1.0) -> Tensor:
r"""Batched :func:`soft_dtw` -> ``(B,)`` for ``cost`` ``(B, n, m)`` (via ``torch.func.vmap``)."""
from torch.func import vmap
def fwd(c: Tensor) -> Tensor:
return soft_dtw(c, beta)
out: Tensor = vmap(fwd)(cost)
return outOther public import paths
omnibias.struct.torch.soft_dtw_batchedomnibias.struct.torch.dtw.soft_dtw_batched
soft_dtw_marginals
omnibias.struct.torch.soft_dtw_marginalsSource ↗soft_dtw_marginals(cost: Tensor, beta: float=1.0) -> TensorClosed-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 | TensorRequired | See the operation contract and implementation below. |
betapositional-or-keyword | floatDefault: 1.0 | See the operation contract and implementation below. |
TensorExecuted regression · soft dtw source and sink marginals are one
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_source_and_sink_marginals_are_one() -> None:
rng = np.random.default_rng(5)
cost = _cost(rng, 4, 4)
e = tdtw.soft_dtw_marginals(torch.tensor(cost), 2.0)
assert abs(float(e[0, 0]) - 1.0) < 1e-9
assert abs(float(e[-1, -1]) - 1.0) < 1e-9python -m pytest "packages/omnibias-struct/tests/test_dtw.py::test_soft_dtw_source_and_sink_marginals_are_one" -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: Tensor, beta: float = 1.0) -> Tensor:
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 = [
torch.stack([torch.exp(-beta * (fwd[i][j] + bwd[i][j] - cost[i, j] - value)) for j in range(m)])
for i in range(n)
]
return torch.stack(rows, dim=0)Other public import paths
omnibias.struct.torch.soft_dtw_marginalsomnibias.struct.torch.dtw.soft_dtw_marginals