API REFERENCE/omnibias-torch

omnibias.torch.line_search

0.5.0rc1Release candidatePython >=3.10Apache-2.0

PyTorch activation derivatives, trainable operator layers, network jets, and direct differential-operator contractions.

import omnibias.torch.line_search

Module contract & mathematical conventions

Exact jet line search (theory 03-12), PyTorch twin.

jet_line_search minimises the directional Taylor model of a scalar phi(s) = loss(params + s * direction). Coefficients come from nested torch.func.grad of the restriction (exact for an omnibias PINN loss whose residual jet is closed form). Model roots, the certified Lagrange radius, and the never-worse verify live in omnibias.core.line_search so the jax twin is bit-identical on the same derivatives.

jet_line_search_on_ray is the input-direction variant: mlp_jet plus compose_jet through a quadratic loss, with no new jet arithmetic.

The driver is eager (host-side root isolation). Do not wrap it in torch.compile and expect the polynomial solve to stay on device.

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 →

Imports & aliases

These names share a definition with another public module. Follow a name for its full contract, parameters, and examples.

Import from this module	Canonical reference
GradientSecant	omnibias.core.GradientSecant
JetLineSearchConfig	omnibias.core.JetLineSearchConfig
LineSearchResult	omnibias.core.LineSearchResult

CONSTANT

Layer

omnibias.torch.line_search.LayerSource ↗

Layer = tuple[Tensor, Tensor | None, ActivationSpec[Tensor] | str | None]

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

Layer = tuple[Tensor, Tensor | None, ActivationSpec[Tensor] | str | None]

CONSTANT

ScalarFn

omnibias.torch.line_search.ScalarFnSource ↗

ScalarFn = Callable[[Tensor], Tensor]

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

ScalarFn = Callable[[Tensor], Tensor]

FUNCTION

directional_derivatives

omnibias.torch.line_search.directional_derivativesSource ↗

directional_derivatives(loss_fn: ScalarFn, params: Tensor, direction: Tensor, order: int) -> list[float]

phi^(k)(0) for k = 0 .. order via nested reverse-mode.

params and direction are flat tensors of the same shape.

Parameters

Name	Type & default	Contract
loss_fnpositional-or-keyword	ScalarFnRequired	See the operation contract and implementation below.
paramspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
directionpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
orderpositional-or-keyword	intRequired	See the operation contract and implementation below.

Returnslist[float]

Explicit validation & failure cases

• When order < 0, raises ValueError(f'order must be >= 0, got {order}').
• When params.shape != direction.shape, raises ValueError(f'params and direction must share shape, got {tuple(params.shape)} vs {tuple(direction.shape)}').

These are explicit source guards. Backend operations may also raise for invalid shapes, dtypes, or unsupported devices.

Executed regression · sweep and named ombu mask

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_sweep_and_named_ombu_mask() -> None:
    torch.set_default_dtype(torch.float64)

    def loss(params: torch.Tensor) -> torch.Tensor:
        p = params.reshape(-1)
        return (p[1] - 2.0) ** 2

    start = torch.zeros(3)
    spec = ombu_bias_block(n_channels=1, k_biases=3, slot=1)
    new, reports = block_exact_sweep(loss, start, spec=spec, exact_quadratic=True)
    assert len(reports) == 1
    assert abs(float(new[1]) - 2.0) <= 1e-12
    assert reports[0].fell_back is False

python -m pytest "packages/omnibias-torch/tests/test_block_exact_search.py::test_sweep_and_named_ombu_mask" -q

Regression example · g6 directional derivatives exp

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_g6_directional_derivatives_exp() -> None:
    jax = pytest.importorskip("jax")
    jax.config.update("jax_enable_x64", True)
    import jax.numpy as jnp
    from omnibias.jax.line_search import directional_derivatives as jax_d
    from omnibias.torch.line_search import directional_derivatives as torch_d

    torch.set_default_dtype(torch.float64)
    p_np = np.array([0.1, -0.2, 0.05], dtype=np.float64)
    d_np = np.array([0.3, 0.4, -0.1], dtype=np.float64)

    def t_loss(p: torch.Tensor) -> torch.Tensor:
        return torch.exp((p**2).sum())

    def j_loss(p: jax.Array) -> jax.Array:
        return jnp.exp(jnp.sum(p**2))

    t_der = torch_d(t_loss, torch.as_tensor(p_np), torch.as_tensor(d_np), 5)
    j_der = jax_d(j_loss, jnp.asarray(p_np), jnp.asarray(d_np), 5)
    worst = max(_ulp_error(a, b) for a, b in zip(t_der, j_der, strict=True))
    assert worst <= 16.0, f"derivative parity worst_ulp={worst}"

python -m pytest tests/test_line_search_parity.py -k test_g6_directional_derivatives_exp -q

Read the implementation

def directional_derivatives(
    loss_fn: ScalarFn,
    params: Tensor,
    direction: Tensor,
    order: int,
) -> list[float]:
    """``phi^(k)(0)`` for ``k = 0 .. order`` via nested reverse-mode.

    ``params`` and ``direction`` are flat tensors of the same shape.
    """
    if order < 0:
        raise ValueError(f"order must be >= 0, got {order}")
    if params.shape != direction.shape:
        raise ValueError(
            f"params and direction must share shape, got {tuple(params.shape)} "
            f"vs {tuple(direction.shape)}"
        )

    def phi(step: Tensor) -> Tensor:
        return loss_fn(params + step * direction)

    s0 = torch.zeros((), dtype=params.dtype, device=params.device)
    derivs = [float(phi(s0))]
    current: Callable[[Tensor], Tensor] = phi
    for _ in range(order):
        current = grad(current)
        derivs.append(float(cast(Tensor, current(s0))))
    return derivs

FUNCTION

jet_line_search

omnibias.torch.jet_line_searchSource ↗

jet_line_search(loss_fn: ScalarFn, params: Tensor, direction: Tensor, *, config: JetLineSearchConfig | None=None, next_derivative_bound: float | None=None, secant: GradientSecant | None=None) -> LineSearchResult

Parameter-direction jet line search with optional certified radius.

secant is the previous steepest step. When it is present the driver may accept an unused inverse eigenvalue of the secant-plane Hessian instead of the single-ray minimizer. verify=True stays the never-worse backstop. The jet is the founding bias collapse.

Parameters

Name	Type & default	Contract
loss_fnpositional-or-keyword	ScalarFnRequired	See the operation contract and implementation below.
paramspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
directionpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
configkeyword-only	JetLineSearchConfig | NoneDefault: None	See the operation contract and implementation below.
next_derivative_boundkeyword-only	float | NoneDefault: None	See the operation contract and implementation below.
secantkeyword-only	GradientSecant | NoneDefault: None	See the operation contract and implementation below.

ReturnsLineSearchResult

Executed regression · sweep and named ombu mask

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_sweep_and_named_ombu_mask() -> None:
    torch.set_default_dtype(torch.float64)

    def loss(params: torch.Tensor) -> torch.Tensor:
        p = params.reshape(-1)
        return (p[1] - 2.0) ** 2

    start = torch.zeros(3)
    spec = ombu_bias_block(n_channels=1, k_biases=3, slot=1)
    new, reports = block_exact_sweep(loss, start, spec=spec, exact_quadratic=True)
    assert len(reports) == 1
    assert abs(float(new[1]) - 2.0) <= 1e-12
    assert reports[0].fell_back is False

python -m pytest "packages/omnibias-torch/tests/test_block_exact_search.py::test_sweep_and_named_ombu_mask" -q

Regression example · quartic recovers spec step

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_quartic_recovers_spec_step() -> None:
    torch.set_default_dtype(torch.float64)
    params = torch.zeros(2, dtype=torch.float64)
    direction = torch.tensor([1.0, 0.0], dtype=torch.float64)
    result = jet_line_search(
        _quartic_loss,
        params,
        direction,
        config=JetLineSearchConfig(order=4, trust_radius=1.0, verify=True),
        next_derivative_bound=0.0,
    )
    assert result.step == pytest.approx(0.409461, abs=5e-5)
    assert result.actual_value is not None
    assert result.actual_value <= 1.0 + 1e-14
    assert result.truncation_certified
    assert not result.fell_back

python -m pytest packages/omnibias-torch/tests/test_line_search.py -k test_quartic_recovers_spec_step -q

Read the implementation

def jet_line_search(
    loss_fn: ScalarFn,
    params: Tensor,
    direction: Tensor,
    *,
    config: JetLineSearchConfig | None = None,
    next_derivative_bound: float | None = None,
    secant: GradientSecant | None = None,
) -> LineSearchResult:
    """Parameter-direction jet line search with optional certified radius.

    ``secant`` is the previous steepest step. When it is present the driver
    may accept an unused inverse eigenvalue of the secant-plane Hessian
    instead of the single-ray minimizer. ``verify=True`` stays the
    never-worse backstop. The jet is the founding bias collapse.
    """
    cfg = config if config is not None else JetLineSearchConfig()
    extra = 1 if (cfg.trust_radius == "certified" and next_derivative_bound is None) else 0
    derivs = directional_derivatives(loss_fn, params, direction, cfg.order + extra)
    bound = next_derivative_bound
    if extra:
        # |phi^(N+1)(0)| is a point value, not a sound max on [0, r].
        bound = None
    params_d = params.detach()
    direction_d = direction.detach()

    def actual(step: float) -> float:
        s = torch.as_tensor(step, dtype=params_d.dtype, device=params_d.device)
        return float(loss_fn(params_d + s * direction_d))

    return run_model_line_search(
        derivs[: cfg.order + 1],
        config=cfg,
        next_derivative_bound=bound,
        actual_fn=actual if cfg.verify else None,
        secant=secant,
    )

Related APIsomnibias.core.JetLineSearchConfig↗omnibias.core.run_model_line_search↗omnibias.torch.line_search.directional_derivatives↗

Other public import paths

• omnibias.torch.jet_line_search
• omnibias.torch.line_search.jet_line_search
• omnibias.torch.optim.jet_line_search

FUNCTION

jet_line_search_on_ray

omnibias.torch.jet_line_search_on_raySource ↗

jet_line_search_on_ray(layers: Sequence[Layer], x0: Tensor, direction: Tensor, target: Tensor, *, config: JetLineSearchConfig | None=None, next_derivative_bound: float | None=None) -> LineSearchResult

Input-ray line search: mlp_jet then quadratic compose_jet.

Parameters

Name	Type & default	Contract
layerspositional-or-keyword	Sequence[Layer]Required	See the operation contract and implementation below.
x0positional-or-keyword	TensorRequired	See the operation contract and implementation below.
directionpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
targetpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
configkeyword-only	JetLineSearchConfig | NoneDefault: None	See the operation contract and implementation below.
next_derivative_boundkeyword-only	float | NoneDefault: None	See the operation contract and implementation below.

ReturnsLineSearchResult

Executed regression · mlp ray uses compose jet

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_mlp_ray_uses_compose_jet() -> None:
    torch.set_default_dtype(torch.float64)
    w = torch.tensor([[0.4, -0.2], [0.1, 0.3]], dtype=torch.float64)
    b = torch.zeros(2, dtype=torch.float64)
    w2 = torch.tensor([[0.5, -0.5]], dtype=torch.float64)
    b2 = torch.zeros(1, dtype=torch.float64)
    layers = [(w, b, "tanh"), (w2, b2, None)]
    x0 = torch.zeros(2, dtype=torch.float64)
    v = torch.tensor([0.3, -0.1], dtype=torch.float64)
    target = torch.tensor([0.0], dtype=torch.float64)
    result = jet_line_search_on_ray(
        layers,
        x0,
        v,
        target,
        config=JetLineSearchConfig(order=4, trust_radius=0.4, verify=True),
    )
    assert result.actual_value is not None
    from omnibias.torch.jet import mlp_jet

    y0 = float(mlp_jet(x0, v, layers, 0)[0, 0])
    start = 0.5 * y0 * y0
    assert result.actual_value <= start + 1e-12

python -m pytest "packages/omnibias-torch/tests/test_line_search.py::test_mlp_ray_uses_compose_jet" -q

Read the implementation

def jet_line_search_on_ray(
    layers: Sequence[Layer],
    x0: Tensor,
    direction: Tensor,
    target: Tensor,
    *,
    config: JetLineSearchConfig | None = None,
    next_derivative_bound: float | None = None,
) -> LineSearchResult:
    """Input-ray line search: ``mlp_jet`` then quadratic ``compose_jet``."""
    cfg = config if config is not None else JetLineSearchConfig()
    y_jet = mlp_jet(x0, direction, layers, cfg.order)
    loss_jet = mse_loss_jet(y_jet, target)
    tower = jet_to_tower(loss_jet)
    derivs = [float(tower[k].reshape(-1)[0]) for k in range(cfg.order + 1)]

    def actual(step: float) -> float:
        moved = x0 + float(step) * direction
        y = mlp_jet(moved, direction, layers, 0)[0]
        tgt = torch.as_tensor(target, dtype=y.dtype, device=y.device)
        residual = y.reshape(-1) - tgt.reshape(-1)
        return float(0.5 * torch.dot(residual, residual))

    return run_model_line_search(
        derivs,
        config=cfg,
        next_derivative_bound=next_derivative_bound,
        actual_fn=actual if cfg.verify else None,
    )

Related APIsomnibias.core.JetLineSearchConfig↗omnibias.core.run_model_line_search↗omnibias.torch.jet_to_tower↗omnibias.torch.mlp_jet↗omnibias.torch.line_search.mse_loss_jet↗

Other public import paths

• omnibias.torch.jet_line_search_on_ray
• omnibias.torch.line_search.jet_line_search_on_ray
• omnibias.torch.optim.jet_line_search_on_ray

FUNCTION

mse_loss_jet

omnibias.torch.line_search.mse_loss_jetSource ↗

mse_loss_jet(output_jet: Tensor, target: Tensor) -> Tensor

Compose 0.5 ||y - t||^2 onto an output jet via compose_jet.

Parameters

Name	Type & default	Contract
output_jetpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
targetpositional-or-keyword	TensorRequired	See the operation contract and implementation below.

ReturnsTensor

Explicit validation & failure cases

• When jet.ndim != 2, raises ValueError(f'output_jet must be (N+1,) or (N+1, C), got {tuple(jet.shape)}').
• When int(tgt.shape[0]) != width, raises ValueError(f'target length {int(tgt.shape[0])} != output width {width}').

These are explicit source guards. Backend operations may also raise for invalid shapes, dtypes, or unsupported devices.

Executed regression · mlp ray uses compose jet

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_mlp_ray_uses_compose_jet() -> None:
    torch.set_default_dtype(torch.float64)
    w = torch.tensor([[0.4, -0.2], [0.1, 0.3]], dtype=torch.float64)
    b = torch.zeros(2, dtype=torch.float64)
    w2 = torch.tensor([[0.5, -0.5]], dtype=torch.float64)
    b2 = torch.zeros(1, dtype=torch.float64)
    layers = [(w, b, "tanh"), (w2, b2, None)]
    x0 = torch.zeros(2, dtype=torch.float64)
    v = torch.tensor([0.3, -0.1], dtype=torch.float64)
    target = torch.tensor([0.0], dtype=torch.float64)
    result = jet_line_search_on_ray(
        layers,
        x0,
        v,
        target,
        config=JetLineSearchConfig(order=4, trust_radius=0.4, verify=True),
    )
    assert result.actual_value is not None
    from omnibias.torch.jet import mlp_jet

    y0 = float(mlp_jet(x0, v, layers, 0)[0, 0])
    start = 0.5 * y0 * y0
    assert result.actual_value <= start + 1e-12

python -m pytest "packages/omnibias-torch/tests/test_line_search.py::test_mlp_ray_uses_compose_jet" -q

Read the implementation

def mse_loss_jet(output_jet: Tensor, target: Tensor) -> Tensor:
    """Compose ``0.5 ||y - t||^2`` onto an output jet via :func:`compose_jet`."""
    jet = torch.as_tensor(output_jet)
    tgt = torch.as_tensor(target, dtype=jet.dtype, device=jet.device)
    if jet.ndim == 1:
        jet = jet.unsqueeze(-1)
    if jet.ndim != 2:
        raise ValueError(f"output_jet must be (N+1,) or (N+1, C), got {tuple(jet.shape)}")
    order = int(jet.shape[0]) - 1
    width = int(jet.shape[1])
    tgt = tgt.reshape(-1)
    if int(tgt.shape[0]) != width:
        raise ValueError(f"target length {int(tgt.shape[0])} != output width {width}")
    acc: Tensor | None = None
    for idx in range(width):
        u = jet[:, idx]
        tower = quadratic_loss_tower(u[0], tgt[idx], order)
        part = compose_jet(u, tower)
        acc = part if acc is None else acc + part
    assert acc is not None
    return acc

Related APIsomnibias.torch.compose_jet↗omnibias.torch.line_search.quadratic_loss_tower↗

FUNCTION

quadratic_loss_tower

omnibias.torch.line_search.quadratic_loss_towerSource ↗

quadratic_loss_tower(u0: Tensor, target: Tensor, order: int) -> Tensor

Derivative tower of 0.5 (u - t)^2 at u0 (orders 0 .. order).

sigma' = u - t, sigma'' = 1, sigma^(k) = 0 for k > 2.

Parameters

Name	Type & default	Contract
u0positional-or-keyword	TensorRequired	See the operation contract and implementation below.
targetpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
orderpositional-or-keyword	intRequired	See the operation contract and implementation below.

ReturnsTensor

Explicit validation & failure cases

• When order < 0, raises ValueError(f'order must be >= 0, got {order}').

These are explicit source guards. Backend operations may also raise for invalid shapes, dtypes, or unsupported devices.

Executed regression · mlp ray uses compose jet

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_mlp_ray_uses_compose_jet() -> None:
    torch.set_default_dtype(torch.float64)
    w = torch.tensor([[0.4, -0.2], [0.1, 0.3]], dtype=torch.float64)
    b = torch.zeros(2, dtype=torch.float64)
    w2 = torch.tensor([[0.5, -0.5]], dtype=torch.float64)
    b2 = torch.zeros(1, dtype=torch.float64)
    layers = [(w, b, "tanh"), (w2, b2, None)]
    x0 = torch.zeros(2, dtype=torch.float64)
    v = torch.tensor([0.3, -0.1], dtype=torch.float64)
    target = torch.tensor([0.0], dtype=torch.float64)
    result = jet_line_search_on_ray(
        layers,
        x0,
        v,
        target,
        config=JetLineSearchConfig(order=4, trust_radius=0.4, verify=True),
    )
    assert result.actual_value is not None
    from omnibias.torch.jet import mlp_jet

    y0 = float(mlp_jet(x0, v, layers, 0)[0, 0])
    start = 0.5 * y0 * y0
    assert result.actual_value <= start + 1e-12

python -m pytest "packages/omnibias-torch/tests/test_line_search.py::test_mlp_ray_uses_compose_jet" -q

Read the implementation

def quadratic_loss_tower(u0: Tensor, target: Tensor, order: int) -> Tensor:
    r"""Derivative tower of ``0.5 (u - t)^2`` at ``u0`` (orders ``0 .. order``).

    ``sigma' = u - t``, ``sigma'' = 1``, ``sigma^(k) = 0`` for ``k > 2``.
    """
    if order < 0:
        raise ValueError(f"order must be >= 0, got {order}")
    rows = [0.5 * (u0 - target) ** 2]
    if order >= 1:
        rows.append(u0 - target)
    if order >= 2:
        rows.append(torch.ones_like(u0))
    rows.extend(torch.zeros_like(u0) for _ in range(max(0, order - 2)))
    return torch.stack(rows[: order + 1], dim=0)

← Package overviewRunnable learning examples →