API REFERENCE/omnibias-struct

omnibias.struct.jax.align

0.1.0a2AlphaPython >=3.10AGPL / commercial

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

import omnibias.struct.jax.align

Module contract & mathematical conventions

Differentiable Needleman-Wunsch sequence alignment (jax).

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

Module source ↗Source revision 9a2d73df0573Running reference examples →

FUNCTION

soft_align

omnibias.struct.jax.soft_alignSource ↗

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

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	ArrayRequired	See the operation contract and implementation below.
gappositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.

ReturnsArray

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 marginals equal grad jax

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_grad_jax() -> None:
    rng = np.random.default_rng(4)
    sub = _blosum_like(4, rng)
    a = np.array([0, 1, 2, 1])
    b = np.array([0, 2, 1])
    beta = 3.0
    sj, gj = jnp.asarray(sub), jnp.asarray(-1.0)
    g_sub = jax.grad(lambda s: jalign.soft_align(a, b, s, gj, beta))(sj)
    g_gap = jax.grad(lambda g: jalign.soft_align(a, b, sj, g, beta))(gj)
    grad_sub, grad_gap = jalign.soft_align_marginals(a, b, sj, gj, beta)
    assert float(jnp.max(jnp.abs(g_sub - grad_sub))) < 1e-9
    assert abs(float(g_gap) - float(grad_gap)) < 1e-9

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

Read the implementation

def soft_align(a: object, b: object, sub: Array, gap: Array, beta: float = 1.0) -> Array:
    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, jnp.asarray(gap, dtype=sub.dtype))
    return -soft_shortest_path(weight, dag, beta)

Related APIsomnibias.struct.jax.soft_shortest_path↗

Other public import paths

• omnibias.struct.jax.soft_align
• omnibias.struct.jax.align.soft_align

FUNCTION

soft_align_batched

omnibias.struct.jax.soft_align_batchedSource ↗

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

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. Maps soft_align over the leading batch axis with jax.vmap -- bit-identical to looping.

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	ArrayRequired	See the operation contract and implementation below.
gappositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.

ReturnsArray

Executed regression · families batched 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.

Full file, imports & fixtures ↗

def test_families_batched_jax() -> None:
    jax = pytest.importorskip("jax")
    jax.config.update("jax_enable_x64", True)
    import jax.numpy as jnp
    from omnibias.struct.jax import (
        soft_align,
        soft_align_batched,
        soft_eisner,
        soft_eisner_batched,
        soft_mas,
        soft_mas_batched,
    )

    rng = np.random.default_rng(4)
    a, b = np.array([0, 1, 2]), np.array([0, 2, 1])
    sub = jnp.asarray(rng.standard_normal((4, 3, 3)))
    gap = jnp.asarray(-0.7)
    b_al = np.asarray(soft_align_batched(a, b, sub, gap, 3.0))
    loop_al = np.asarray([float(soft_align(a, b, sub[i], gap, 3.0)) for i in range(4)])
    assert np.max(np.abs(b_al - loop_al)) < ATOL

    score = jnp.asarray(rng.standard_normal((4, 3, 4)))
    b_mas = np.asarray(soft_mas_batched(score, 2.0))
    loop_mas = np.asarray([float(soft_mas(score[i], 2.0)) for i in range(4)])
    assert np.max(np.abs(b_mas - loop_mas)) < ATOL

    arc = jnp.asarray(rng.standard_normal((4, 4, 4)))
    b_ei = np.asarray(soft_eisner_batched(arc, 2.0))
    loop_ei = np.asarray([float(soft_eisner(arc[i], 2.0)) for i in range(4)])
    assert np.max(np.abs(b_ei - loop_ei)) < ATOL

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

Regression example · soft align batched matches loop 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.

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 -k test_soft_align_batched_matches_loop_torch -q

Read the implementation

def soft_align_batched(a: object, b: object, sub: Array, gap: Array, beta: float = 1.0) -> Array:
    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. Maps :func:`soft_align`
    over the leading batch axis with ``jax.vmap`` -- bit-identical to looping.
    """
    import jax

    gap_a = jnp.asarray(gap, dtype=sub.dtype)
    gap_dim = 0 if gap_a.ndim == 1 else None

    def fwd(s: Array, g: Array) -> Array:
        return soft_align(a, b, s, g, beta)

    out: Array = jax.vmap(fwd, in_axes=(0, gap_dim))(sub, gap_a)
    return out

Related APIsomnibias.struct.jax.soft_align↗

Other public import paths

• omnibias.struct.jax.soft_align_batched
• omnibias.struct.jax.align.soft_align_batched

FUNCTION

soft_align_marginals

omnibias.struct.jax.soft_align_marginalsSource ↗

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

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 jax.grad 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	ArrayRequired	See the operation contract and implementation below.
gappositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.

Returnstuple[Array, Array]

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

Full file, imports & fixtures ↗

def test_soft_align_marginals_equal_grad_jax() -> None:
    rng = np.random.default_rng(4)
    sub = _blosum_like(4, rng)
    a = np.array([0, 1, 2, 1])
    b = np.array([0, 2, 1])
    beta = 3.0
    sj, gj = jnp.asarray(sub), jnp.asarray(-1.0)
    g_sub = jax.grad(lambda s: jalign.soft_align(a, b, s, gj, beta))(sj)
    g_gap = jax.grad(lambda g: jalign.soft_align(a, b, sj, g, beta))(gj)
    grad_sub, grad_gap = jalign.soft_align_marginals(a, b, sj, gj, beta)
    assert float(jnp.max(jnp.abs(g_sub - grad_sub))) < 1e-9
    assert abs(float(g_gap) - float(grad_gap)) < 1e-9

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

Read the implementation

def soft_align_marginals(
    a: object, b: object, sub: Array, gap: Array, beta: float = 1.0
) -> tuple[Array, Array]:
    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 ``jax.grad`` of :func:`soft_align`.
    """
    gap_a = jnp.asarray(gap, dtype=sub.dtype)
    weight, dag, labels, _ = _assemble(a, b, sub, gap_a)
    xi = soft_shortest_path_marginals(weight, dag, beta)
    ai = np.asarray(a, dtype=int).reshape(-1)
    bj = np.asarray(b, dtype=int).reshape(-1)
    grad_sub = jnp.zeros_like(sub)
    grad_gap = jnp.zeros((), dtype=sub.dtype)
    for (u, v), (kind, i, j) in labels.items():
        if kind == "sub":
            grad_sub = grad_sub.at[int(ai[i]), int(bj[j])].add(xi[u, v])
        else:
            grad_gap = grad_gap + xi[u, v]
    return grad_sub, grad_gap

Related APIsomnibias.struct.jax.soft_shortest_path_marginals↗

Other public import paths

• omnibias.struct.jax.soft_align_marginals
• omnibias.struct.jax.align.soft_align_marginals

FUNCTION

soft_gotoh

omnibias.struct.jax.soft_gotohSource ↗

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

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	ArrayRequired	See the operation contract and implementation below.
gap_openpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
gap_extendpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.

ReturnsArray

Executed regression · alignment parity bit identical

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_parity_bit_identical(kind: str) -> None:
    pytest.importorskip("torch")
    pytest.importorskip("jax")
    assert probe_align_parity(kind, 0, BETA).metrics["value_abs_diff"] < 1e-9

python -m pytest "packages/omnibias-struct/tests/test_refinement_families.py::test_alignment_parity_bit_identical[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: Array, gap_open: Array, gap_extend: Array, beta: float = 1.0
) -> Array:
    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, jnp.asarray(gap_open, dtype=sub.dtype), jnp.asarray(gap_extend, dtype=sub.dtype)
    )
    value: Array = -soft_shortest_path(weight, dag, beta)
    return value

Related APIsomnibias.struct.jax.soft_shortest_path↗

Other public import paths

• omnibias.struct.jax.soft_gotoh
• omnibias.struct.jax.align.soft_gotoh

FUNCTION

soft_gotoh_marginals

omnibias.struct.jax.soft_gotoh_marginalsSource ↗

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

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 jax.grad.

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	ArrayRequired	See the operation contract and implementation below.
gap_openpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
gap_extendpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.

Returnstuple[Array, Array, Array]

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: Array, gap_open: Array, gap_extend: Array, beta: float = 1.0
) -> tuple[Array, Array, Array]:
    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 ``jax.grad``.
    """
    open_a = jnp.asarray(gap_open, dtype=sub.dtype)
    extend_a = jnp.asarray(gap_extend, dtype=sub.dtype)
    weight, dag, labels = _assemble_gotoh(a, b, sub, open_a, extend_a)
    xi = soft_shortest_path_marginals(weight, dag, beta)
    ai = np.asarray(a, dtype=int).reshape(-1)
    bj = np.asarray(b, dtype=int).reshape(-1)
    grad_sub = jnp.zeros_like(sub)
    grad_open = jnp.zeros((), dtype=sub.dtype)
    grad_extend = jnp.zeros((), dtype=sub.dtype)
    for (u, v), (kind, i, j) in labels.items():
        if kind == "sub":
            grad_sub = grad_sub.at[int(ai[i]), int(bj[j])].add(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.jax.soft_shortest_path_marginals↗

Other public import paths

• omnibias.struct.jax.soft_gotoh_marginals
• omnibias.struct.jax.align.soft_gotoh_marginals

FUNCTION

soft_local_align

omnibias.struct.jax.soft_local_alignSource ↗

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

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	ArrayRequired	See the operation contract and implementation below.
gappositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.

ReturnsArray

Executed regression · alignment parity bit identical

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_parity_bit_identical(kind: str) -> None:
    pytest.importorskip("torch")
    pytest.importorskip("jax")
    assert probe_align_parity(kind, 0, BETA).metrics["value_abs_diff"] < 1e-9

python -m pytest "packages/omnibias-struct/tests/test_refinement_families.py::test_alignment_parity_bit_identical[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: Array, gap: Array, beta: float = 1.0) -> Array:
    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, jnp.asarray(gap, dtype=sub.dtype))
    value: Array = -soft_shortest_path(weight, dag, beta)
    return value

Related APIsomnibias.struct.jax.soft_shortest_path↗

Other public import paths

• omnibias.struct.jax.soft_local_align
• omnibias.struct.jax.align.soft_local_align

FUNCTION

soft_local_align_marginals

omnibias.struct.jax.soft_local_align_marginalsSource ↗

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

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	ArrayRequired	See the operation contract and implementation below.
gappositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.

Returnstuple[Array, Array]

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: Array, gap: Array, beta: float = 1.0
) -> tuple[Array, Array]:
    r"""Closed-form parameter gradients ``(d/d sub, d/d gap)`` of :func:`soft_local_align`."""
    gap_a = jnp.asarray(gap, dtype=sub.dtype)
    weight, dag, labels = _assemble_local(a, b, sub, gap_a)
    xi = soft_shortest_path_marginals(weight, dag, beta)
    ai = np.asarray(a, dtype=int).reshape(-1)
    bj = np.asarray(b, dtype=int).reshape(-1)
    grad_sub = jnp.zeros_like(sub)
    grad_gap = jnp.zeros((), dtype=sub.dtype)
    for (u, v), (kind, i, j) in labels.items():
        if kind == "sub":
            grad_sub = grad_sub.at[int(ai[i]), int(bj[j])].add(xi[u, v])
        elif kind == "gap":
            grad_gap = grad_gap + xi[u, v]
    return grad_sub, grad_gap

Related APIsomnibias.struct.jax.soft_shortest_path_marginals↗

Other public import paths

• omnibias.struct.jax.soft_local_align_marginals
• omnibias.struct.jax.align.soft_local_align_marginals

← Package overviewRunnable learning examples →