API REFERENCE/omnibias-jax

omnibias.jax.line_search

0.5.0rc1Release candidatePython >=3.10Apache-2.0

JAX activation derivatives, functional network jets, direct contractions, and parameter-differentiable numerical models.

import omnibias.jax.line_search

Module contract & mathematical conventions

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

Bit-identical partner of omnibias.torch.line_search on the same directional derivatives. The polynomial solve is host-side (eager); do not jit the driver. Nested jax.grad of the scalar restriction is the jet; mlp_jet / compose_jet handle the input-ray path.

Enable 64-bit JAX before the first array if you need torch parity (omnibias.jax.precision).

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.jax.line_search.LayerSource ↗

Layer = tuple[Array, Array | None, ActivationSpec[Array] | 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[Array, Array | None, ActivationSpec[Array] | str | None]

CONSTANT

ScalarFn

omnibias.jax.line_search.ScalarFnSource ↗

ScalarFn = Callable[[Any], Array]

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[[Any], Array]

FUNCTION

directional_derivatives

omnibias.jax.line_search.directional_derivativesSource ↗

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

phi^(k)(0) for k = 0 .. order via nested jax.grad.

Parameters

Name	Type & default	Contract
loss_fnpositional-or-keyword	ScalarFnRequired	See the operation contract and implementation below.
paramspositional-or-keyword	AnyRequired	See the operation contract and implementation below.
directionpositional-or-keyword	AnyRequired	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}').

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

Executed regression · g3 never worse pathological

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_g3_never_worse_pathological() -> None:
    def sharp_bowl(p: jax.Array) -> jax.Array:
        s = p[0]
        return 1.0 - 2.0 * s + 50.0 * s**2

    result = jet_line_search(
        sharp_bowl,
        jnp.zeros(1),
        jnp.ones(1),
        config=JetLineSearchConfig(order=1, trust_radius=1.0, verify=True),
    )
    assert result.actual_value is not None
    assert result.actual_value <= 1.0 + 1e-14
    assert result.fell_back

python -m pytest "packages/omnibias-jax/tests/test_line_search.py::test_g3_never_worse_pathological" -q

Regression example · g1 jet matches closed form orders 0 to 6

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_g1_jet_matches_closed_form_orders_0_to_6() -> None:
    def loss(p: jax.Array) -> jax.Array:
        return jnp.exp(jnp.sum(p))

    params = jnp.zeros(3)
    direction = jnp.ones(3)
    derivs = directional_derivatives(loss, params, direction, 6)
    for k, val in enumerate(derivs):
        expected = 3.0**k
        rel = abs(val - expected) / max(abs(expected), 1.0)
        assert rel <= 1e-10, f"order {k}: {val} vs {expected} rel={rel}"

python -m pytest packages/omnibias-jax/tests/test_line_search.py -k test_g1_jet_matches_closed_form_orders_0_to_6 -q

Read the implementation

def directional_derivatives(
    loss_fn: ScalarFn,
    params: Any,
    direction: Any,
    order: int,
) -> list[float]:
    """``phi^(k)(0)`` for ``k = 0 .. order`` via nested ``jax.grad``."""
    if order < 0:
        raise ValueError(f"order must be >= 0, got {order}")
    theta, direction_flat, unravel = _flat_pair(params, direction)

    def phi(step: Array) -> Array:
        return loss_fn(unravel(theta + step * direction_flat))

    s0 = jnp.zeros((), dtype=theta.dtype)
    derivs = [float(phi(s0))]
    current: Callable[[Array], Array] = phi
    for _ in range(order):
        current = jax.grad(current)
        derivs.append(float(current(s0)))
    return derivs

FUNCTION

jet_line_search

omnibias.jax.jet_line_searchSource ↗

jet_line_search(loss_fn: ScalarFn, params: Any, direction: Any, *, 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	AnyRequired	See the operation contract and implementation below.
directionpositional-or-keyword	AnyRequired	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 · g3 never worse pathological

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_g3_never_worse_pathological() -> None:
    def sharp_bowl(p: jax.Array) -> jax.Array:
        s = p[0]
        return 1.0 - 2.0 * s + 50.0 * s**2

    result = jet_line_search(
        sharp_bowl,
        jnp.zeros(1),
        jnp.ones(1),
        config=JetLineSearchConfig(order=1, trust_radius=1.0, verify=True),
    )
    assert result.actual_value is not None
    assert result.actual_value <= 1.0 + 1e-14
    assert result.fell_back

python -m pytest "packages/omnibias-jax/tests/test_line_search.py::test_g3_never_worse_pathological" -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:
    params = jnp.zeros(2)
    direction = jnp.array([1.0, 0.0])
    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-jax/tests/test_line_search.py -k test_quartic_recovers_spec_step -q

Read the implementation

def jet_line_search(
    loss_fn: ScalarFn,
    params: Any,
    direction: Any,
    *,
    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:
        bound = None
    theta, direction_flat, unravel = _flat_pair(params, direction)

    def actual(step: float) -> float:
        s = jnp.asarray(step, dtype=theta.dtype)
        return float(loss_fn(unravel(theta + s * direction_flat)))

    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.jax.line_search.directional_derivatives↗

Other public import paths

• omnibias.jax.jet_line_search
• omnibias.jax.line_search.jet_line_search
• omnibias.jax.optim.jet_line_search

FUNCTION

jet_line_search_on_ray

omnibias.jax.jet_line_search_on_raySource ↗

jet_line_search_on_ray(layers: Sequence[Layer], x0: Array, direction: Array, target: Array, *, 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	ArrayRequired	See the operation contract and implementation below.
directionpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
targetpositional-or-keyword	ArrayRequired	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

Read the implementation

def jet_line_search_on_ray(
    layers: Sequence[Layer],
    x0: Array,
    direction: Array,
    target: Array,
    *,
    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(jnp.reshape(tower[k], (-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 = jnp.asarray(target, dtype=y.dtype)
        residual = jnp.reshape(y, (-1,)) - jnp.reshape(tgt, (-1,))
        return float(0.5 * jnp.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.jax.jet_to_tower↗omnibias.jax.mlp_jet↗omnibias.jax.line_search.mse_loss_jet↗

Other public import paths

• omnibias.jax.jet_line_search_on_ray
• omnibias.jax.line_search.jet_line_search_on_ray
• omnibias.jax.optim.jet_line_search_on_ray

FUNCTION

mse_loss_jet

omnibias.jax.line_search.mse_loss_jetSource ↗

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

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

Parameters

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

ReturnsArray

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.

Read the implementation

def mse_loss_jet(output_jet: Array, target: Array) -> Array:
    """Compose ``0.5 ||y - t||^2`` onto an output jet via :func:`compose_jet`."""
    jet = jnp.asarray(output_jet)
    tgt = jnp.asarray(target, dtype=jet.dtype)
    if jet.ndim == 1:
        jet = jet[:, None]
    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: Array | 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.jax.compose_jet↗omnibias.jax.line_search.quadratic_loss_tower↗

FUNCTION

quadratic_loss_tower

omnibias.jax.line_search.quadratic_loss_towerSource ↗

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

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

Parameters

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

ReturnsArray

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.

Read the implementation

def quadratic_loss_tower(u0: Array, target: Array, order: int) -> Array:
    r"""Derivative tower of ``0.5 (u - t)^2`` at ``u0`` (orders ``0 .. order``)."""
    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(jnp.ones_like(u0))
    rows.extend(jnp.zeros_like(u0) for _ in range(max(0, order - 2)))
    return jnp.stack(rows[: order + 1], axis=0)

← Package overviewRunnable learning examples →