API REFERENCE/omnibias-struct

omnibias.struct.torch.align

0.1.0a2AlphaPython >=3.10AGPL / commercial

Differentiable dynamic programs for sequences and parses, structured distributions, and finite-temperature bounds.

import omnibias.struct.torch.align

Module contract & mathematical conventions

Differentiable Needleman-Wunsch sequence alignment (torch).

Bit-identical twin of omnibias.struct.jax.align (float64). Global alignment is a longest-path DAG, so this *reuses the shared shortest-path substrate*: assemble a weight matrix W[u, v] = -score(edge) (negate so longest-path = shortest-cost), differentiable in the substitution matrix sub and gap penalty gap, then call omnibias.struct.torch.soft_shortest_path. The beta -> inf softmax anneals to the NW optimum; the delta -> 0 tower gives the closed-form parameter gradients (soft_align_marginals), which equal autograd. 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.

Module source ↗Source revision 9a2d73df0573Running reference examples →

FUNCTION

soft_align

omnibias.struct.torch.soft_alignSource ↗

soft_align(a: object, b: object, sub: Tensor, gap: Tensor, beta: float=1.0) -> Tensor

Soft global-alignment score beta^-1 log sum_paths exp(beta score) of a vs b.

sub is a (K, K) substitution-score matrix and gap a scalar gap penalty (both learnable). Differentiable in sub and gap; -> Needleman-Wunsch optimum as beta -> inf.

Parameters

Name	Type & default	Contract
apositional-or-keyword	objectRequired	See the operation contract and implementation below.
bpositional-or-keyword	objectRequired	See the operation contract and implementation below.
subpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
gappositional-or-keyword	TensorRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.

ReturnsTensor

Executed regression · soft align torch jax 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.

Full file, imports & fixtures ↗

def test_soft_align_torch_jax_parity() -> None:
    rng = np.random.default_rng(5)
    sub = _blosum_like(5, rng)
    a = np.array([0, 3, 2, 1])
    b = np.array([0, 2, 4, 1])
    for beta in (1.0, 8.0):
        v_t = float(talign.soft_align(a, b, torch.tensor(sub), torch.tensor(-0.7), beta))
        v_j = float(jalign.soft_align(a, b, jnp.asarray(sub), jnp.asarray(-0.7), beta))
        assert abs(v_t - v_j) < 1e-9

python -m pytest "packages/omnibias-struct/tests/test_align.py::test_soft_align_torch_jax_parity" -q

Regression example · soft align matches flat softmax and anneals

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.

Full file, imports & fixtures ↗

def test_soft_align_matches_flat_softmax_and_anneals() -> None:
    rng = np.random.default_rng(1)
    sub = _blosum_like(4, rng)
    gap = -1.0
    a = np.array([0, 1, 2, 1])
    b = np.array([0, 2, 1])
    st, gt = torch.tensor(sub), torch.tensor(gap)
    hard = hard_align(a, b, sub, gap)
    for beta in (0.5, 1.0, 4.0, 16.0):
        soft = float(talign.soft_align(a, b, st, gt, beta))
        assert abs(soft - brute_force_soft_align(a, b, sub, gap, beta)) < 1e-9
        assert soft >= hard - 1e-9  # softmax >= max

python -m pytest packages/omnibias-struct/tests/test_align.py -k test_soft_align_matches_flat_softmax_and_anneals -q

Read the implementation

def soft_align(a: object, b: object, sub: Tensor, gap: Tensor, beta: float = 1.0) -> Tensor:
    r"""Soft global-alignment score ``beta^-1 log sum_paths exp(beta score)`` of ``a`` vs ``b``.

    ``sub`` is a ``(K, K)`` substitution-score matrix and ``gap`` a scalar gap penalty (both
    learnable). Differentiable in ``sub`` and ``gap``; ``-> Needleman-Wunsch optimum`` as
    ``beta -> inf``.
    """
    weight, dag, _, _ = _assemble(a, b, sub, torch.as_tensor(gap, dtype=sub.dtype))
    return -soft_shortest_path(weight, dag, beta)

Related APIsomnibias.struct.torch.soft_shortest_path↗

Other public import paths

• omnibias.struct.torch.soft_align
• omnibias.struct.torch.align.soft_align

FUNCTION

soft_align_batched

omnibias.struct.torch.soft_align_batchedSource ↗

soft_align_batched(a: object, b: object, sub: Tensor, gap: Tensor, beta: float=1.0) -> Tensor

Batched soft_align -> (B,) for sub (B, K, K) (shared sequences a, b).

gap may be a shared scalar or a per-example (B,) vector. The alignment DAG is identical across the batch (same a / b), so this assembles the batched (B, n, n) edge-cost tensor once and hands the reduction to the vmap-based omnibias.struct.torch.soft_shortest_path_batched -- bit-identical to looping soft_align.

Parameters

Name	Type & default	Contract
apositional-or-keyword	objectRequired	See the operation contract and implementation below.
bpositional-or-keyword	objectRequired	See the operation contract and implementation below.
subpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
gappositional-or-keyword	TensorRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.

ReturnsTensor

Executed regression · soft align batched matches loop 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.

Full file, imports & fixtures ↗

def test_soft_align_batched_matches_loop_torch() -> None:
    torch = pytest.importorskip("torch")
    torch.set_default_dtype(torch.float64)
    from omnibias.struct.torch import soft_align, soft_align_batched

    rng = np.random.default_rng(0)
    a, b = np.array([0, 1, 2, 1]), np.array([0, 2, 1])
    sub = torch.tensor(rng.standard_normal((5, 3, 3)))
    gap = torch.tensor(rng.standard_normal(5))  # per-example
    batched = soft_align_batched(a, b, sub, gap, 3.0)
    loop = torch.stack([soft_align(a, b, sub[i], gap[i], 3.0) for i in range(5)])
    assert batched.shape == (5,)
    assert torch.max(torch.abs(batched - loop)).item() < ATOL
    # shared scalar gap also works
    shared = soft_align_batched(a, b, sub, torch.tensor(-1.0), 3.0)
    loop2 = torch.stack([soft_align(a, b, sub[i], torch.tensor(-1.0), 3.0) for i in range(5)])
    assert torch.max(torch.abs(shared - loop2)).item() < ATOL

python -m pytest "packages/omnibias-struct/tests/test_batch_families.py::test_soft_align_batched_matches_loop_torch" -q

Read the implementation

def soft_align_batched(a: object, b: object, sub: Tensor, gap: Tensor, beta: float = 1.0) -> Tensor:
    r"""Batched :func:`soft_align` -> ``(B,)`` for ``sub`` ``(B, K, K)`` (shared sequences ``a``, ``b``).

    ``gap`` may be a shared scalar or a per-example ``(B,)`` vector. The alignment DAG is
    identical across the batch (same ``a`` / ``b``), so this assembles the batched
    ``(B, n, n)`` edge-cost tensor once and hands the reduction to the ``vmap``-based
    :func:`omnibias.struct.torch.soft_shortest_path_batched` -- bit-identical to looping
    :func:`soft_align`.
    """
    ai = np.asarray(a, dtype=int).reshape(-1)
    bj = np.asarray(b, dtype=int).reshape(-1)
    batch = sub.shape[0]
    gap_t = torch.as_tensor(gap, dtype=sub.dtype)
    gap_vec = gap_t.expand(batch) if gap_t.ndim == 0 else gap_t
    lattice = AlignmentLattice(ai.shape[0], bj.shape[0])
    dag, labels = lattice.build_dag()
    weight = torch.zeros((batch, lattice.num_nodes, lattice.num_nodes), dtype=sub.dtype, device=sub.device)
    for (u, v), (kind, i, j) in labels.items():
        weight[:, u, v] = -sub[:, int(ai[i]), int(bj[j])] if kind == "sub" else -gap_vec
    value: Tensor = -soft_shortest_path_batched(weight, dag, beta)
    return value

Related APIsomnibias.struct.AlignmentLattice↗omnibias.struct.torch.soft_shortest_path_batched↗

Other public import paths

• omnibias.struct.torch.soft_align_batched
• omnibias.struct.torch.align.soft_align_batched

FUNCTION

soft_align_marginals

omnibias.struct.torch.soft_align_marginalsSource ↗

soft_align_marginals(a: object, b: object, sub: Tensor, gap: Tensor, beta: float=1.0) -> tuple[Tensor, Tensor]

Closed-form parameter gradients (d/d sub, d/d gap) of soft_align.

Forward-backward edge marginals xi from the tower softmax, folded back onto the parameters: grad_sub[p, q] is the expected number of (p, q) substitutions and grad_gap the expected number of gaps. Equal to autograd of soft_align.

Parameters

Name	Type & default	Contract
apositional-or-keyword	objectRequired	See the operation contract and implementation below.
bpositional-or-keyword	objectRequired	See the operation contract and implementation below.
subpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
gappositional-or-keyword	TensorRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.

Returnstuple[Tensor, Tensor]

Executed regression · soft align 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.

Full file, imports & fixtures ↗

def test_soft_align_marginals_equal_autograd_torch() -> None:
    rng = np.random.default_rng(3)
    sub = _blosum_like(4, rng)
    a = np.array([0, 1, 2, 1])
    b = np.array([0, 2, 1])
    beta = 3.0
    st = torch.tensor(sub, requires_grad=True)
    gt = torch.tensor(-1.0, requires_grad=True)
    talign.soft_align(a, b, st, gt, beta).backward()
    grad_sub, grad_gap = talign.soft_align_marginals(a, b, st.detach(), gt.detach(), beta)
    assert torch.max(torch.abs(st.grad - grad_sub)).item() < 1e-9
    assert abs(float(gt.grad) - float(grad_gap)) < 1e-9

python -m pytest "packages/omnibias-struct/tests/test_align.py::test_soft_align_marginals_equal_autograd_torch" -q

Read the implementation

def soft_align_marginals(
    a: object, b: object, sub: Tensor, gap: Tensor, beta: float = 1.0
) -> tuple[Tensor, Tensor]:
    r"""Closed-form parameter gradients ``(d/d sub, d/d gap)`` of :func:`soft_align`.

    Forward-backward edge marginals ``xi`` from the tower softmax, folded back onto the
    parameters: ``grad_sub[p, q]`` is the expected number of ``(p, q)`` substitutions and
    ``grad_gap`` the expected number of gaps. Equal to ``autograd`` of :func:`soft_align`.
    """
    gap_t = torch.as_tensor(gap, dtype=sub.dtype)
    weight, dag, labels, _ = _assemble(a, b, sub, gap_t)
    xi = soft_shortest_path_marginals(weight, dag, beta).detach()
    ai = np.asarray(a, dtype=int).reshape(-1)
    bj = np.asarray(b, dtype=int).reshape(-1)
    grad_sub = torch.zeros_like(sub)
    grad_gap = torch.zeros((), dtype=sub.dtype, device=sub.device)
    for (u, v), (kind, i, j) in labels.items():
        if kind == "sub":
            grad_sub[int(ai[i]), int(bj[j])] += xi[u, v]
        else:
            grad_gap = grad_gap + xi[u, v]
    return grad_sub, grad_gap

Related APIsomnibias.struct.torch.soft_shortest_path_marginals↗

Other public import paths

• omnibias.struct.torch.soft_align_marginals
• omnibias.struct.torch.align.soft_align_marginals

FUNCTION

soft_gotoh

omnibias.struct.torch.soft_gotohSource ↗

soft_gotoh(a: object, b: object, sub: Tensor, gap_open: Tensor, gap_extend: Tensor, beta: float=1.0) -> Tensor

Soft Gotoh affine-gap alignment score beta^-1 log sum_paths exp(beta score).

A gap of length L scores gap_open + L * gap_extend (both learnable). Differentiable in sub / gap_open / gap_extend; -> the Gotoh optimum as beta -> inf.

Parameters

Name	Type & default	Contract
apositional-or-keyword	objectRequired	See the operation contract and implementation below.
bpositional-or-keyword	objectRequired	See the operation contract and implementation below.
subpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
gap_openpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
gap_extendpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.

ReturnsTensor

Executed regression · alignment soft ge hard

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.

Full file, imports & fixtures ↗

def test_alignment_soft_ge_hard(kind: str) -> None:
    pytest.importorskip("torch")
    assert probe_align_gap(kind, 0, BETA, "torch").ok

python -m pytest "packages/omnibias-struct/tests/test_refinement_families.py::test_alignment_soft_ge_hard[affine]" -q

Regression example · soft gotoh 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.

Full file, imports & fixtures ↗

def test_soft_gotoh_torch_jax_parity(seed: int) -> None:
    rng = np.random.default_rng(seed)
    sub = _sub(5, rng)
    a, b = _seqs(rng, 5)
    for beta in (1.0, 8.0):
        v_t = float(talign.soft_gotoh(a, b, torch.tensor(sub), torch.tensor(-1.2), torch.tensor(-0.3), beta))
        v_j = float(
            jalign.soft_gotoh(a, b, jnp.asarray(sub), jnp.asarray(-1.2), jnp.asarray(-0.3), beta)
        )
        assert abs(v_t - v_j) < 1e-9

python -m pytest packages/omnibias-struct/tests/test_align_local_affine.py -k test_soft_gotoh_torch_jax_parity -q

Read the implementation

def soft_gotoh(
    a: object, b: object, sub: Tensor, gap_open: Tensor, gap_extend: Tensor, beta: float = 1.0
) -> Tensor:
    r"""Soft Gotoh affine-gap alignment score ``beta^-1 log sum_paths exp(beta score)``.

    A gap of length ``L`` scores ``gap_open + L * gap_extend`` (both learnable). Differentiable
    in ``sub`` / ``gap_open`` / ``gap_extend``; ``-> the Gotoh optimum`` as ``beta -> inf``.
    """
    weight, dag, _ = _assemble_gotoh(
        a, b, sub, torch.as_tensor(gap_open, dtype=sub.dtype), torch.as_tensor(gap_extend, dtype=sub.dtype)
    )
    value: Tensor = -soft_shortest_path(weight, dag, beta)
    return value

Related APIsomnibias.struct.torch.soft_shortest_path↗

Other public import paths

• omnibias.struct.torch.soft_gotoh
• omnibias.struct.torch.align.soft_gotoh

FUNCTION

soft_gotoh_marginals

omnibias.struct.torch.soft_gotoh_marginalsSource ↗

soft_gotoh_marginals(a: object, b: object, sub: Tensor, gap_open: Tensor, gap_extend: Tensor, beta: float=1.0) -> tuple[Tensor, Tensor, Tensor]

Closed-form gradients (d/d sub, d/d gap_open, d/d gap_extend) of soft_gotoh.

Each gap-open edge scores gap_open + gap_extend (so it feeds both gradients) and each gap-extend edge scores gap_extend; folding the forward-backward edge marginals back gives the expected substitution / open / (open + extend) counts, equal to autograd.

Parameters

Name	Type & default	Contract
apositional-or-keyword	objectRequired	See the operation contract and implementation below.
bpositional-or-keyword	objectRequired	See the operation contract and implementation below.
subpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
gap_openpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
gap_extendpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.

Returnstuple[Tensor, Tensor, Tensor]

Executed regression · soft gotoh marginals equal autograd

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.

Full file, imports & fixtures ↗

def test_soft_gotoh_marginals_equal_autograd() -> None:
    rng = np.random.default_rng(3)
    sub = _sub(4, rng)
    a = np.array([0, 1, 2, 1])
    b = np.array([0, 2, 1])
    beta = 3.0
    # torch closed-form == autograd
    st = torch.tensor(sub, requires_grad=True)
    got = torch.tensor(-1.5, requires_grad=True)
    get = torch.tensor(-0.4, requires_grad=True)
    talign.soft_gotoh(a, b, st, got, get, beta).backward()
    g_sub, g_open, g_ext = talign.soft_gotoh_marginals(a, b, st.detach(), got.detach(), get.detach(), beta)
    assert torch.max(torch.abs(st.grad - g_sub)).item() < 1e-9
    assert abs(float(got.grad) - float(g_open)) < 1e-9
    assert abs(float(get.grad) - float(g_ext)) < 1e-9
    # jax closed-form == jax.grad
    sj, oj, ej = jnp.asarray(sub), jnp.asarray(-1.5), jnp.asarray(-0.4)
    jg_sub = jax.grad(lambda s: jalign.soft_gotoh(a, b, s, oj, ej, beta))(sj)
    jg_open = jax.grad(lambda o: jalign.soft_gotoh(a, b, sj, o, ej, beta))(oj)
    jg_ext = jax.grad(lambda e: jalign.soft_gotoh(a, b, sj, oj, e, beta))(ej)
    cf_sub, cf_open, cf_ext = jalign.soft_gotoh_marginals(a, b, sj, oj, ej, beta)
    assert float(jnp.max(jnp.abs(jg_sub - cf_sub))) < 1e-9
    assert abs(float(jg_open) - float(cf_open)) < 1e-9
    assert abs(float(jg_ext) - float(cf_ext)) < 1e-9

python -m pytest "packages/omnibias-struct/tests/test_align_local_affine.py::test_soft_gotoh_marginals_equal_autograd" -q

Read the implementation

def soft_gotoh_marginals(
    a: object, b: object, sub: Tensor, gap_open: Tensor, gap_extend: Tensor, beta: float = 1.0
) -> tuple[Tensor, Tensor, Tensor]:
    r"""Closed-form gradients ``(d/d sub, d/d gap_open, d/d gap_extend)`` of :func:`soft_gotoh`.

    Each gap-open edge scores ``gap_open + gap_extend`` (so it feeds both gradients) and each
    gap-extend edge scores ``gap_extend``; folding the forward-backward edge marginals back
    gives the expected substitution / open / (open + extend) counts, equal to ``autograd``.
    """
    open_t = torch.as_tensor(gap_open, dtype=sub.dtype)
    extend_t = torch.as_tensor(gap_extend, dtype=sub.dtype)
    weight, dag, labels = _assemble_gotoh(a, b, sub, open_t, extend_t)
    xi = soft_shortest_path_marginals(weight, dag, beta).detach()
    ai = np.asarray(a, dtype=int).reshape(-1)
    bj = np.asarray(b, dtype=int).reshape(-1)
    grad_sub = torch.zeros_like(sub)
    grad_open = torch.zeros((), dtype=sub.dtype, device=sub.device)
    grad_extend = torch.zeros((), dtype=sub.dtype, device=sub.device)
    for (u, v), (kind, i, j) in labels.items():
        if kind == "sub":
            grad_sub[int(ai[i]), int(bj[j])] += xi[u, v]
        elif kind == "open":
            grad_open = grad_open + xi[u, v]
            grad_extend = grad_extend + xi[u, v]
        elif kind == "extend":
            grad_extend = grad_extend + xi[u, v]
    return grad_sub, grad_open, grad_extend

Related APIsomnibias.struct.torch.soft_shortest_path_marginals↗

Other public import paths

• omnibias.struct.torch.soft_gotoh_marginals
• omnibias.struct.torch.align.soft_gotoh_marginals

FUNCTION

soft_local_align

omnibias.struct.torch.soft_local_alignSource ↗

soft_local_align(a: object, b: object, sub: Tensor, gap: Tensor, beta: float=1.0) -> Tensor

Soft Smith-Waterman local-alignment score beta^-1 log sum_paths exp(beta score).

Free start / end 0-edges on the shared shortest-path substrate, so -> the Smith-Waterman local optimum as beta -> inf (with the empty alignment as the score-0 floor). Differentiable in the substitution matrix sub and gap penalty gap.

Parameters

Name	Type & default	Contract
apositional-or-keyword	objectRequired	See the operation contract and implementation below.
bpositional-or-keyword	objectRequired	See the operation contract and implementation below.
subpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
gappositional-or-keyword	TensorRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.

ReturnsTensor

Executed regression · alignment soft ge hard

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.

Full file, imports & fixtures ↗

def test_alignment_soft_ge_hard(kind: str) -> None:
    pytest.importorskip("torch")
    assert probe_align_gap(kind, 0, BETA, "torch").ok

python -m pytest "packages/omnibias-struct/tests/test_refinement_families.py::test_alignment_soft_ge_hard[local]" -q

Regression example · soft local 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.

Full file, imports & fixtures ↗

def test_soft_local_torch_jax_parity(seed: int) -> None:
    rng = np.random.default_rng(seed)
    sub = _sub(5, rng)
    a, b = _seqs(rng, 5)
    for beta in (1.0, 8.0):
        v_t = float(talign.soft_local_align(a, b, torch.tensor(sub), torch.tensor(-0.7), beta))
        v_j = float(jalign.soft_local_align(a, b, jnp.asarray(sub), jnp.asarray(-0.7), beta))
        assert abs(v_t - v_j) < 1e-9

python -m pytest packages/omnibias-struct/tests/test_align_local_affine.py -k test_soft_local_torch_jax_parity -q

Read the implementation

def soft_local_align(a: object, b: object, sub: Tensor, gap: Tensor, beta: float = 1.0) -> Tensor:
    r"""Soft Smith-Waterman local-alignment score ``beta^-1 log sum_paths exp(beta score)``.

    Free start / end 0-edges on the shared shortest-path substrate, so ``-> the
    Smith-Waterman local optimum`` as ``beta -> inf`` (with the empty alignment as the
    score-``0`` floor). Differentiable in the substitution matrix ``sub`` and gap penalty
    ``gap``.
    """
    weight, dag, _ = _assemble_local(a, b, sub, torch.as_tensor(gap, dtype=sub.dtype))
    value: Tensor = -soft_shortest_path(weight, dag, beta)
    return value

Related APIsomnibias.struct.torch.soft_shortest_path↗

Other public import paths

• omnibias.struct.torch.soft_local_align
• omnibias.struct.torch.align.soft_local_align

FUNCTION

soft_local_align_marginals

omnibias.struct.torch.soft_local_align_marginalsSource ↗

soft_local_align_marginals(a: object, b: object, sub: Tensor, gap: Tensor, beta: float=1.0) -> tuple[Tensor, Tensor]

Closed-form parameter gradients (d/d sub, d/d gap) of soft_local_align.

Parameters

Name	Type & default	Contract
apositional-or-keyword	objectRequired	See the operation contract and implementation below.
bpositional-or-keyword	objectRequired	See the operation contract and implementation below.
subpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
gappositional-or-keyword	TensorRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.

Returnstuple[Tensor, Tensor]

Executed regression · soft local marginals equal autograd

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.

Full file, imports & fixtures ↗

def test_soft_local_marginals_equal_autograd() -> None:
    rng = np.random.default_rng(3)
    sub = _sub(4, rng)
    a = np.array([0, 1, 2, 1])
    b = np.array([0, 2, 1])
    beta = 3.0
    # torch closed-form == autograd
    st = torch.tensor(sub, requires_grad=True)
    gt = torch.tensor(-1.0, requires_grad=True)
    talign.soft_local_align(a, b, st, gt, beta).backward()
    g_sub, g_gap = talign.soft_local_align_marginals(a, b, st.detach(), gt.detach(), beta)
    assert torch.max(torch.abs(st.grad - g_sub)).item() < 1e-9
    assert abs(float(gt.grad) - float(g_gap)) < 1e-9
    # jax closed-form == jax.grad
    sj, gj = jnp.asarray(sub), jnp.asarray(-1.0)
    jg_sub = jax.grad(lambda s: jalign.soft_local_align(a, b, s, gj, beta))(sj)
    jg_gap = jax.grad(lambda g: jalign.soft_local_align(a, b, sj, g, beta))(gj)
    cf_sub, cf_gap = jalign.soft_local_align_marginals(a, b, sj, gj, beta)
    assert float(jnp.max(jnp.abs(jg_sub - cf_sub))) < 1e-9
    assert abs(float(jg_gap) - float(cf_gap)) < 1e-9

python -m pytest "packages/omnibias-struct/tests/test_align_local_affine.py::test_soft_local_marginals_equal_autograd" -q

Read the implementation

def soft_local_align_marginals(
    a: object, b: object, sub: Tensor, gap: Tensor, beta: float = 1.0
) -> tuple[Tensor, Tensor]:
    r"""Closed-form parameter gradients ``(d/d sub, d/d gap)`` of :func:`soft_local_align`."""
    gap_t = torch.as_tensor(gap, dtype=sub.dtype)
    weight, dag, labels = _assemble_local(a, b, sub, gap_t)
    xi = soft_shortest_path_marginals(weight, dag, beta).detach()
    ai = np.asarray(a, dtype=int).reshape(-1)
    bj = np.asarray(b, dtype=int).reshape(-1)
    grad_sub = torch.zeros_like(sub)
    grad_gap = torch.zeros((), dtype=sub.dtype, device=sub.device)
    for (u, v), (kind, i, j) in labels.items():
        if kind == "sub":
            grad_sub[int(ai[i]), int(bj[j])] += xi[u, v]
        elif kind == "gap":
            grad_gap = grad_gap + xi[u, v]
    return grad_sub, grad_gap

Related APIsomnibias.struct.torch.soft_shortest_path_marginals↗

Other public import paths

• omnibias.struct.torch.soft_local_align_marginals
• omnibias.struct.torch.align.soft_local_align_marginals

← Package overviewRunnable learning examples →