API REFERENCE/omnibias-torch

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

Module contract & mathematical conventions

Block / coordinate exact search (theory 08-07), PyTorch 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

Do not wrap the driver in torch.compile.

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.torch.block_directionSource ↗

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

Unit direction supported on a named or boolean block.

Parameters

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

ReturnsTensor

Executed regression · empty mask raises

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_empty_mask_raises() -> None:
    torch.set_default_dtype(torch.float64)
    with pytest.raises(ValueError, match="empty"):
        block_exact_search(_section5_loss, torch.zeros(2), mask=(False, False))

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

Read the implementation

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

Related APIsomnibias.core.unit_direction_from_mask↗

Other public import paths

• omnibias.torch.block_direction
• omnibias.torch.optim.block_direction
• omnibias.torch.optim_block_search.block_direction

FUNCTION

block_exact_search

omnibias.torch.block_exact_searchSource ↗

block_exact_search(loss_fn: ScalarFn, params: Tensor, *, spec: BlockSpec | None=None, mask: Tensor | Sequence[bool] | None=None, direction: Tensor | None=None, config: JetLineSearchConfig | None=None, exact_quadratic: bool=False, next_derivative_bound: float | None=None) -> tuple[Tensor, 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	TensorRequired	See the operation contract and implementation below.
speckeyword-only	BlockSpec | NoneDefault: None	See the operation contract and implementation below.
maskkeyword-only	Tensor | Sequence[bool] | NoneDefault: None	See the operation contract and implementation below.
directionkeyword-only	Tensor | 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[Tensor, 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 · empty mask raises

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_empty_mask_raises() -> None:
    torch.set_default_dtype(torch.float64)
    with pytest.raises(ValueError, match="empty"):
        block_exact_search(_section5_loss, torch.zeros(2), mask=(False, False))

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

Regression example · g2 exact quadratic

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:
    torch.set_default_dtype(torch.float64)
    v = torch.tensor([0.0, 0.0])
    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
    assert abs(result.model_value) <= 1e-12

python -m pytest packages/omnibias-torch/tests/test_block_exact_search.py -k test_g2_exact_quadratic -q

Read the implementation

def block_exact_search(
    loss_fn: ScalarFn,
    params: Tensor,
    *,
    spec: BlockSpec | None = None,
    mask: Tensor | Sequence[bool] | None = None,
    direction: Tensor | None = None,
    config: JetLineSearchConfig | None = None,
    exact_quadratic: bool = False,
    next_derivative_bound: float | None = None,
) -> tuple[Tensor, LineSearchResult]:
    """03-12 line search along one structured block; never-worse by default."""
    p = params.reshape(-1)
    flags = _as_mask(p, spec=spec, mask=mask)
    if direction is None:
        g = grad(loss_fn)(p)
        probe = (-g).reshape(-1)
        d = block_direction(p, spec=BlockSpec(kind="mask"), mask=flags, probe=probe)
    else:
        raw = direction.reshape(-1)
        if int(raw.shape[0]) != int(p.shape[0]):
            raise ValueError("direction must match flat params")
        restricted = torch.where(
            torch.as_tensor(flags, dtype=torch.bool, device=p.device),
            raw,
            torch.zeros_like(raw),
        )
        nrm = float(torch.linalg.vector_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 torch.as_tensor(new, dtype=p.dtype, device=p.device), result

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

Other public import paths

• omnibias.torch.block_exact_search
• omnibias.torch.optim.block_exact_search
• omnibias.torch.optim_block_search.block_exact_search

FUNCTION

block_exact_sweep

omnibias.torch.block_exact_sweepSource ↗

block_exact_sweep(loss_fn: ScalarFn, params: Tensor, *, spec: BlockSpec | None=None, mask: Tensor | Sequence[bool] | None=None, config: JetLineSearchConfig | None=None, exact_quadratic: bool=False, next_derivative_bound: float | None=None) -> tuple[Tensor, 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	TensorRequired	See the operation contract and implementation below.
speckeyword-only	BlockSpec | NoneDefault: None	See the operation contract and implementation below.
maskkeyword-only	Tensor | 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[Tensor, list[LineSearchResult]]

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

Read the implementation

def block_exact_sweep(
    loss_fn: ScalarFn,
    params: Tensor,
    *,
    spec: BlockSpec | None = None,
    mask: Tensor | Sequence[bool] | None = None,
    config: JetLineSearchConfig | None = None,
    exact_quadratic: bool = False,
    next_derivative_bound: float | None = None,
) -> tuple[Tensor, list[LineSearchResult]]:
    """One Gauss–Seidel pass: search each singleton coordinate of the block."""
    p = params.reshape(-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.torch.block_exact_search↗

Other public import paths

• omnibias.torch.block_exact_sweep
• omnibias.torch.optim.block_exact_sweep
• omnibias.torch.optim_block_search.block_exact_sweep

← Package overviewRunnable learning examples →