API REFERENCE/omnibias-jax

omnibias.jax.optim_block_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.optim_block_search

Module contract & mathematical conventions

Block / coordinate exact search (theory 08-07), JAX twin.

A sparse block direction is handed to the 03-12 jet line search (verify=True never-worse). Last-layer least squares is exactly quadratic. This is a coordinate / block sweep, not a global solver

loss_fn must be jit-compatible. Enable x64 before the first array when matching the torch twin.

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
BlockSpec	omnibias.core.BlockSpec
arrangement_w_block	omnibias.core.arrangement_w_block
last_linear_block	omnibias.core.last_linear_block
ombu_bias_block	omnibias.core.ombu_bias_block

FUNCTION

block_direction

omnibias.jax.block_directionSource ↗

block_direction(params: Array, *, spec: BlockSpec | None=None, mask: Array | Sequence[bool] | None=None, probe: Array | Sequence[float] | None=None) -> Array

Unit direction supported on a named or boolean block.

Parameters

Name	Type & default	Contract
paramspositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
speckeyword-only	BlockSpec | NoneDefault: None	See the operation contract and implementation below.
maskkeyword-only	Array | Sequence[bool] | NoneDefault: None	See the operation contract and implementation below.
probekeyword-only	Array | Sequence[float] | NoneDefault: None	See the operation contract and implementation below.

ReturnsArray

Executed regression · g2 exact quadratic

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_g2_exact_quadratic() -> None:
    v = jnp.array([0.0, 0.0], dtype=jnp.float64)
    new, result = block_exact_search(
        _section5_loss,
        v,
        mask=(True, False),
        exact_quadratic=True,
    )
    assert result.fell_back is False
    assert abs(result.step - 1.0) <= 1e-12
    assert abs(float(new[0]) - 1.0) <= 1e-12
    assert abs(float(new[1])) <= 1e-12
    assert result.actual_value is not None
    assert abs(result.actual_value) <= 1e-12

python -m pytest "packages/omnibias-jax/tests/test_block_exact_search.py::test_g2_exact_quadratic" -q

Read the implementation

def block_direction(
    params: Array,
    *,
    spec: BlockSpec | None = None,
    mask: Array | Sequence[bool] | None = None,
    probe: Array | Sequence[float] | None = None,
) -> Array:
    """Unit direction supported on a named or boolean block."""
    p = jnp.reshape(params, (-1,))
    flags = _as_mask(p, spec=spec, mask=mask)
    probe_list: Sequence[float] | None
    if probe is None:
        probe_list = None
    else:
        probe_list = [float(v) for v in jnp.reshape(jnp.asarray(probe), (-1,))]
    direction = unit_direction_from_mask(flags, probe_list)
    return jnp.asarray(direction, dtype=p.dtype)

Related APIsomnibias.core.unit_direction_from_mask↗

Other public import paths

• omnibias.jax.block_direction
• omnibias.jax.optim.block_direction
• omnibias.jax.optim_block_search.block_direction

FUNCTION

block_exact_search

omnibias.jax.block_exact_searchSource ↗

block_exact_search(loss_fn: ScalarFn, params: Array, *, spec: BlockSpec | None=None, mask: Array | Sequence[bool] | None=None, direction: Array | None=None, config: JetLineSearchConfig | None=None, exact_quadratic: bool=False, next_derivative_bound: float | None=None) -> tuple[Array, LineSearchResult]

03-12 line search along one structured block; never-worse by default.

Parameters

Name	Type & default	Contract
loss_fnpositional-or-keyword	ScalarFnRequired	See the operation contract and implementation below.
paramspositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
speckeyword-only	BlockSpec | NoneDefault: None	See the operation contract and implementation below.
maskkeyword-only	Array | Sequence[bool] | NoneDefault: None	See the operation contract and implementation below.
directionkeyword-only	Array | NoneDefault: None	See the operation contract and implementation below.
configkeyword-only	JetLineSearchConfig | NoneDefault: None	See the operation contract and implementation below.
exact_quadratickeyword-only	boolDefault: False	See the operation contract and implementation below.
next_derivative_boundkeyword-only	float | NoneDefault: None	See the operation contract and implementation below.

Returnstuple[Array, LineSearchResult]

Explicit validation & failure cases

• When int(raw.shape[0]) != int(p.shape[0]), raises ValueError('direction must match flat params').

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

Executed regression · g2 exact quadratic

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_g2_exact_quadratic() -> None:
    v = jnp.array([0.0, 0.0], dtype=jnp.float64)
    new, result = block_exact_search(
        _section5_loss,
        v,
        mask=(True, False),
        exact_quadratic=True,
    )
    assert result.fell_back is False
    assert abs(result.step - 1.0) <= 1e-12
    assert abs(float(new[0]) - 1.0) <= 1e-12
    assert abs(float(new[1])) <= 1e-12
    assert result.actual_value is not None
    assert abs(result.actual_value) <= 1e-12

python -m pytest "packages/omnibias-jax/tests/test_block_exact_search.py::test_g2_exact_quadratic" -q

Regression example · last linear block decreases

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_last_linear_block_decreases() -> None:
    hidden = jnp.array([[1.0, 0.0], [0.0, 1.0]], dtype=jnp.float64)
    target = jnp.array([1.0, -1.0], dtype=jnp.float64)

    def loss(params: Array) -> Array:
        r = hidden @ jnp.reshape(params, (-1,)) - target
        return 0.5 * jnp.dot(r, r)

    v0 = jnp.zeros(2, dtype=jnp.float64)
    new, result = block_exact_search(
        loss, v0, spec=last_linear_block(2), exact_quadratic=True
    )
    assert float(loss(new)) <= float(loss(v0))
    assert result.fell_back is False

python -m pytest packages/omnibias-jax/tests/test_block_exact_search.py -k test_last_linear_block_decreases -q

Read the implementation

def block_exact_search(
    loss_fn: ScalarFn,
    params: Array,
    *,
    spec: BlockSpec | None = None,
    mask: Array | Sequence[bool] | None = None,
    direction: Array | None = None,
    config: JetLineSearchConfig | None = None,
    exact_quadratic: bool = False,
    next_derivative_bound: float | None = None,
) -> tuple[Array, LineSearchResult]:
    """03-12 line search along one structured block; never-worse by default."""
    p = jnp.reshape(params, (-1,))
    flags = _as_mask(p, spec=spec, mask=mask)
    if direction is None:
        g = jax.grad(loss_fn)(p)
        probe = jnp.reshape(-g, (-1,))
        d = block_direction(p, spec=BlockSpec(kind="mask"), mask=flags, probe=probe)
    else:
        raw = jnp.reshape(direction, (-1,))
        if int(raw.shape[0]) != int(p.shape[0]):
            raise ValueError("direction must match flat params")
        flag_arr = jnp.asarray(flags)
        restricted = jnp.where(flag_arr, raw, jnp.zeros_like(raw))
        nrm = float(jnp.linalg.norm(restricted))
        if nrm == 0.0:
            d = block_direction(p, spec=BlockSpec(kind="mask"), mask=flags)
        else:
            d = restricted / nrm
    cfg = config if config is not None else default_block_config(
        exact_quadratic=exact_quadratic
    )
    bound = 0.0 if exact_quadratic and next_derivative_bound is None else next_derivative_bound
    result = jet_line_search(
        loss_fn, p, d, config=cfg, next_derivative_bound=bound
    )
    new = apply_block_step(
        [float(v) for v in p], [float(v) for v in d], result.step
    )
    return jnp.asarray(new, dtype=p.dtype), result

Related APIsomnibias.core.BlockSpec↗omnibias.core.apply_block_step↗omnibias.core.default_block_config↗omnibias.jax.jet_line_search↗omnibias.jax.block_direction↗

Other public import paths

• omnibias.jax.block_exact_search
• omnibias.jax.optim.block_exact_search
• omnibias.jax.optim_block_search.block_exact_search

FUNCTION

block_exact_sweep

omnibias.jax.block_exact_sweepSource ↗

block_exact_sweep(loss_fn: ScalarFn, params: Array, *, spec: BlockSpec | None=None, mask: Array | Sequence[bool] | None=None, config: JetLineSearchConfig | None=None, exact_quadratic: bool=False, next_derivative_bound: float | None=None) -> tuple[Array, list[LineSearchResult]]

One Gauss–Seidel pass: search each singleton coordinate of the block.

Parameters

Name	Type & default	Contract
loss_fnpositional-or-keyword	ScalarFnRequired	See the operation contract and implementation below.
paramspositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
speckeyword-only	BlockSpec | NoneDefault: None	See the operation contract and implementation below.
maskkeyword-only	Array | Sequence[bool] | NoneDefault: None	See the operation contract and implementation below.
configkeyword-only	JetLineSearchConfig | NoneDefault: None	See the operation contract and implementation below.
exact_quadratickeyword-only	boolDefault: False	See the operation contract and implementation below.
next_derivative_boundkeyword-only	float | NoneDefault: None	See the operation contract and implementation below.

Returnstuple[Array, list[LineSearchResult]]

Read the implementation

def block_exact_sweep(
    loss_fn: ScalarFn,
    params: Array,
    *,
    spec: BlockSpec | None = None,
    mask: Array | Sequence[bool] | None = None,
    config: JetLineSearchConfig | None = None,
    exact_quadratic: bool = False,
    next_derivative_bound: float | None = None,
) -> tuple[Array, list[LineSearchResult]]:
    """One Gauss–Seidel pass: search each singleton coordinate of the block."""
    p = jnp.reshape(params, (-1,))
    flags = _as_mask(p, spec=spec, mask=mask)
    reports: list[LineSearchResult] = []
    for i, active in enumerate(flags):
        if not active:
            continue
        coord = [False] * len(flags)
        coord[i] = True
        p, report = block_exact_search(
            loss_fn,
            p,
            spec=BlockSpec(kind="mask"),
            mask=coord,
            config=config,
            exact_quadratic=exact_quadratic,
            next_derivative_bound=next_derivative_bound,
        )
        reports.append(report)
    return p, reports

Related APIsomnibias.core.BlockSpec↗omnibias.jax.block_exact_search↗

Other public import paths

• omnibias.jax.block_exact_sweep
• omnibias.jax.optim.block_exact_sweep
• omnibias.jax.optim_block_search.block_exact_sweep

← Package overviewRunnable learning examples →