API REFERENCE/omnibias-core

omnibias.core.block_search

0.5.0rc1Release candidatePython >=3.10Apache-2.0

Coefficient recurrences, activation contracts, jet algebra, and verified arithmetic shared by the numerical backends.

import omnibias.core.block_search

Module contract & mathematical conventions

Block / coordinate exact search algebra (theory 08-07).

A structured block (last linear layer, one OMBU bias, one arrangement normal) is a sparse direction. The step along that direction is spec 03-12: a Taylor model of phi(s), optional certified radius, and verify=True never-worse. Last-layer least squares is exactly quadratic, so order 2 with remainder bound 0 is exact.

This is a coordinate / block sweep, not a global solver and with explicit local scope. Bias collapse (delta -> 0) supplies the tower. Arrangement beta is caller-owned; this step is at fixed beta.

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 →

CONSTANT

BlockKind

omnibias.core.BlockKindSource ↗

BlockKind = Literal['mask', 'last_linear', 'ombu_bias', 'arrangement_W']

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

BlockKind = Literal["mask", "last_linear", "ombu_bias", "arrangement_W"]

Other public import paths

• omnibias.core.BlockKind
• omnibias.core.block_search.BlockKind

CLASS

BlockSpec

omnibias.core.BlockSpecSource ↗

class BlockSpec

Dataclass constructor fields; inherited fields, if any, are declared by the base class. A field(...) expression describes a default factory or field options, not a literal argument default.

Which coordinates of a flat parameter vector may move.

Parameters

Name	Type & default	Contract
kinddataclass field	BlockKindDefault: 'mask'	See the operation contract and implementation below.
widthdataclass field	int | NoneDefault: None	See the operation contract and implementation below.
indexdataclass field	intDefault: 0	See the operation contract and implementation below.
slotdataclass field	intDefault: 0	See the operation contract and implementation below.
n_channelsdataclass field	int | NoneDefault: None	See the operation contract and implementation below.
k_biasesdataclass field	int | NoneDefault: None	See the operation contract and implementation below.
n_featuresdataclass field	int | NoneDefault: None	See the operation contract and implementation below.
n_hyperplanesdataclass field	int | NoneDefault: None	See the operation contract and implementation below.
offsetdataclass field	intDefault: 0	See the operation contract and implementation below.

Executed regression · last linear mask is the tail

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_last_linear_mask_is_the_tail() -> None:
    spec = last_linear_block(2)
    assert resolve_block_mask(5, spec) == (False, False, False, True, True)

python -m pytest "packages/omnibias-core/tests/test_block_search.py::test_last_linear_mask_is_the_tail" -q

Regression example · empty mask rejected

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_rejected() -> None:
    with pytest.raises(ValueError, match="empty"):
        unit_direction_from_mask((False, False))
    with pytest.raises(ValueError):
        BlockSpec(kind="last_linear", width=0)

python -m pytest packages/omnibias-core/tests/test_block_search.py -k test_empty_mask_rejected -q

Fields & methods

kindfield

kind: BlockKind = 'mask'

Declared BlockKind field on BlockSpec. Default expression: 'mask'.

View source ↗

widthfield

width: int | None = None

Declared int | None field on BlockSpec. Default expression: None.

View source ↗

indexfield

index: int = 0

Declared int field on BlockSpec. Default expression: 0.

View source ↗

slotfield

slot: int = 0

Declared int field on BlockSpec. Default expression: 0.

View source ↗

n_channelsfield

n_channels: int | None = None

Declared int | None field on BlockSpec. Default expression: None.

View source ↗

k_biasesfield

k_biases: int | None = None

Declared int | None field on BlockSpec. Default expression: None.

View source ↗

n_featuresfield

n_features: int | None = None

Declared int | None field on BlockSpec. Default expression: None.

View source ↗

n_hyperplanesfield

n_hyperplanes: int | None = None

Declared int | None field on BlockSpec. Default expression: None.

View source ↗

offsetfield

offset: int = 0

Declared int field on BlockSpec. Default expression: 0.

View source ↗

__post_init__method

__post_init__(self) -> None

See the class contract above and the exact method implementation below.

Returns: None

Validation & failure cases

• When self.kind not in ('mask', 'last_linear', 'ombu_bias', 'arrangement_W'), raises ValueError(f'unknown block kind {self.kind!r}').
• When self.index < 0, raises ValueError(f'index must be >= 0, got {self.index}').
• When self.slot < 0, raises ValueError(f'slot must be >= 0, got {self.slot}').
• When self.offset < 0, raises ValueError(f'offset must be >= 0, got {self.offset}').
• When self.width is None or int(self.width) < 1, raises ValueError('last_linear requires width >= 1').
• When self.n_channels is None or int(self.n_channels) < 1, raises ValueError('ombu_bias requires n_channels >= 1').
• When self.k_biases is None or int(self.k_biases) < 1, raises ValueError('ombu_bias requires k_biases >= 1').
• When self.index >= int(self.n_channels), raises ValueError(f'channel index {self.index} >= n_channels {self.n_channels}').
• When self.slot >= int(self.k_biases), raises ValueError(f'slot {self.slot} >= k_biases {self.k_biases}').
• When self.n_features is None or int(self.n_features) < 1, raises ValueError('arrangement_W requires n_features >= 1').
• When self.n_hyperplanes is None or int(self.n_hyperplanes) < 1, raises ValueError('arrangement_W requires n_hyperplanes >= 1').
• When self.index >= int(self.n_hyperplanes), raises ValueError(f'row index {self.index} >= n_hyperplanes {self.n_hyperplanes}').

Executed regression · last linear mask is the tail

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_last_linear_mask_is_the_tail() -> None:
    spec = last_linear_block(2)
    assert resolve_block_mask(5, spec) == (False, False, False, True, True)

python -m pytest "packages/omnibias-core/tests/test_block_search.py::test_last_linear_mask_is_the_tail" -q

Method implementation

    def __post_init__(self) -> None:
        if self.kind not in ("mask", "last_linear", "ombu_bias", "arrangement_W"):
            raise ValueError(f"unknown block kind {self.kind!r}")
        if self.index < 0:
            raise ValueError(f"index must be >= 0, got {self.index}")
        if self.slot < 0:
            raise ValueError(f"slot must be >= 0, got {self.slot}")
        if self.offset < 0:
            raise ValueError(f"offset must be >= 0, got {self.offset}")
        if self.kind == "last_linear":
            if self.width is None or int(self.width) < 1:
                raise ValueError("last_linear requires width >= 1")
        if self.kind == "ombu_bias":
            if self.n_channels is None or int(self.n_channels) < 1:
                raise ValueError("ombu_bias requires n_channels >= 1")
            if self.k_biases is None or int(self.k_biases) < 1:
                raise ValueError("ombu_bias requires k_biases >= 1")
            if self.index >= int(self.n_channels):
                raise ValueError(
                    f"channel index {self.index} >= n_channels {self.n_channels}"
                )
            if self.slot >= int(self.k_biases):
                raise ValueError(f"slot {self.slot} >= k_biases {self.k_biases}")
        if self.kind == "arrangement_W":
            if self.n_features is None or int(self.n_features) < 1:
                raise ValueError("arrangement_W requires n_features >= 1")
            if self.n_hyperplanes is None or int(self.n_hyperplanes) < 1:
                raise ValueError("arrangement_W requires n_hyperplanes >= 1")
            if self.index >= int(self.n_hyperplanes):
                raise ValueError(
                    f"row index {self.index} >= n_hyperplanes {self.n_hyperplanes}"
                )

View source ↗

Read the implementation

class BlockSpec:
    """Which coordinates of a flat parameter vector may move."""

    kind: BlockKind = "mask"
    width: int | None = None
    index: int = 0
    slot: int = 0
    n_channels: int | None = None
    k_biases: int | None = None
    n_features: int | None = None
    n_hyperplanes: int | None = None
    offset: int = 0

    def __post_init__(self) -> None:
        if self.kind not in ("mask", "last_linear", "ombu_bias", "arrangement_W"):
            raise ValueError(f"unknown block kind {self.kind!r}")
        if self.index < 0:
            raise ValueError(f"index must be >= 0, got {self.index}")
        if self.slot < 0:
            raise ValueError(f"slot must be >= 0, got {self.slot}")
        if self.offset < 0:
            raise ValueError(f"offset must be >= 0, got {self.offset}")
        if self.kind == "last_linear":
            if self.width is None or int(self.width) < 1:
                raise ValueError("last_linear requires width >= 1")
        if self.kind == "ombu_bias":
            if self.n_channels is None or int(self.n_channels) < 1:
                raise ValueError("ombu_bias requires n_channels >= 1")
            if self.k_biases is None or int(self.k_biases) < 1:
                raise ValueError("ombu_bias requires k_biases >= 1")
            if self.index >= int(self.n_channels):
                raise ValueError(
                    f"channel index {self.index} >= n_channels {self.n_channels}"
                )
            if self.slot >= int(self.k_biases):
                raise ValueError(f"slot {self.slot} >= k_biases {self.k_biases}")
        if self.kind == "arrangement_W":
            if self.n_features is None or int(self.n_features) < 1:
                raise ValueError("arrangement_W requires n_features >= 1")
            if self.n_hyperplanes is None or int(self.n_hyperplanes) < 1:
                raise ValueError("arrangement_W requires n_hyperplanes >= 1")
            if self.index >= int(self.n_hyperplanes):
                raise ValueError(
                    f"row index {self.index} >= n_hyperplanes {self.n_hyperplanes}"
                )

Other public import paths

• omnibias.core.BlockSpec
• omnibias.jax.BlockSpec
• omnibias.torch.BlockSpec
• omnibias.core.block_search.BlockSpec
• omnibias.jax.optim.BlockSpec
• omnibias.jax.optim_block_search.BlockSpec
• omnibias.torch.optim.BlockSpec
• omnibias.torch.optim_block_search.BlockSpec

FUNCTION

apply_block_step

omnibias.core.apply_block_stepSource ↗

apply_block_step(params: Sequence[float], direction: Sequence[float], step: float) -> tuple[float, ...]

params + step * direction as a flat tuple.

Parameters

Name	Type & default	Contract
paramspositional-or-keyword	Sequence[float]Required	See the operation contract and implementation below.
directionpositional-or-keyword	Sequence[float]Required	See the operation contract and implementation below.
steppositional-or-keyword	floatRequired	See the operation contract and implementation below.

Returnstuple[float, ...]

Explicit validation & failure cases

• When len(params) != len(direction), raises ValueError('params and direction must have the same length').
• When not math.isfinite(float(step)), raises ValueError(f'step must be finite, got {step!r}').

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

Executed regression · apply block step

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_apply_block_step() -> None:
    assert apply_block_step((0.0, 0.0), (1.0, 0.0), 1.0) == (1.0, 0.0)

python -m pytest "packages/omnibias-core/tests/test_block_search.py::test_apply_block_step" -q

Read the implementation

def apply_block_step(
    params: Sequence[float],
    direction: Sequence[float],
    step: float,
) -> tuple[float, ...]:
    """``params + step * direction`` as a flat tuple."""
    if len(params) != len(direction):
        raise ValueError("params and direction must have the same length")
    if not math.isfinite(float(step)):
        raise ValueError(f"step must be finite, got {step!r}")
    return tuple(float(p) + float(step) * float(d) for p, d in zip(params, direction, strict=True))

Other public import paths

• omnibias.core.apply_block_step
• omnibias.core.block_search.apply_block_step

FUNCTION

arrangement_w_block

omnibias.core.arrangement_w_blockSource ↗

arrangement_w_block(*, n_hyperplanes: int, n_features: int, row: int=0, offset: int=0) -> BlockSpec

One arrangement normal (one row of W), row-major.

Parameters

Name	Type & default	Contract
n_hyperplaneskeyword-only	intRequired	See the operation contract and implementation below.
n_featureskeyword-only	intRequired	See the operation contract and implementation below.
rowkeyword-only	intDefault: 0	See the operation contract and implementation below.
offsetkeyword-only	intDefault: 0	See the operation contract and implementation below.

ReturnsBlockSpec

Executed regression · arrangement w is one row

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_arrangement_w_is_one_row() -> None:
    spec = arrangement_w_block(n_hyperplanes=2, n_features=3, row=1)
    flags = resolve_block_mask(6, spec)
    assert flags == (False, False, False, True, True, True)

python -m pytest "packages/omnibias-core/tests/test_block_search.py::test_arrangement_w_is_one_row" -q

Regression example · last linear mask is the tail

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_mask_is_the_tail() -> None:
    spec = last_linear_block(2)
    assert resolve_block_mask(5, spec) == (False, False, False, True, True)

python -m pytest packages/omnibias-core/tests/test_block_search.py -k test_last_linear_mask_is_the_tail -q

Read the implementation

def arrangement_w_block(
    *,
    n_hyperplanes: int,
    n_features: int,
    row: int = 0,
    offset: int = 0,
) -> BlockSpec:
    """One arrangement normal (one row of ``W``), row-major."""
    return BlockSpec(
        kind="arrangement_W",
        n_hyperplanes=int(n_hyperplanes),
        n_features=int(n_features),
        index=int(row),
        offset=int(offset),
    )

Related APIsomnibias.core.BlockSpec↗

Other public import paths

• omnibias.core.arrangement_w_block
• omnibias.jax.arrangement_w_block
• omnibias.torch.arrangement_w_block
• omnibias.core.block_search.arrangement_w_block
• omnibias.jax.optim.arrangement_w_block
• omnibias.jax.optim_block_search.arrangement_w_block
• omnibias.torch.optim.arrangement_w_block
• omnibias.torch.optim_block_search.arrangement_w_block

FUNCTION

default_block_config

omnibias.core.default_block_configSource ↗

default_block_config(*, exact_quadratic: bool=False) -> JetLineSearchConfig

03-12 config. Last-layer LS uses order 2 and a larger explicit radius.

Parameters

Name	Type & default	Contract
exact_quadratickeyword-only	boolDefault: False	See the operation contract and implementation below.

ReturnsJetLineSearchConfig

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 default_block_config(*, exact_quadratic: bool = False) -> JetLineSearchConfig:
    """03-12 config. Last-layer LS uses order 2 and a larger explicit radius."""
    if exact_quadratic:
        return JetLineSearchConfig(
            order=2,
            trust_radius=8.0,
            verify=True,
            max_step=8.0,
        )
    return JetLineSearchConfig(verify=True)

Related APIsomnibias.core.JetLineSearchConfig↗

Other public import paths

• omnibias.core.default_block_config
• omnibias.core.block_search.default_block_config

FUNCTION

last_linear_block

omnibias.core.last_linear_blockSource ↗

last_linear_block(width: int, *, offset: int=0) -> BlockSpec

Last-linear slice: width entries starting at offset from the tail.

Parameters

Name	Type & default	Contract
widthpositional-or-keyword	intRequired	See the operation contract and implementation below.
offsetkeyword-only	intDefault: 0	See the operation contract and implementation below.

ReturnsBlockSpec

Executed regression · last linear mask is the tail

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_last_linear_mask_is_the_tail() -> None:
    spec = last_linear_block(2)
    assert resolve_block_mask(5, spec) == (False, False, False, True, True)

python -m pytest "packages/omnibias-core/tests/test_block_search.py::test_last_linear_mask_is_the_tail" -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 last_linear_block(width: int, *, offset: int = 0) -> BlockSpec:
    """Last-linear slice: ``width`` entries starting at ``offset`` from the tail."""
    return BlockSpec(kind="last_linear", width=int(width), offset=int(offset))

Related APIsomnibias.core.BlockSpec↗

Other public import paths

• omnibias.core.last_linear_block
• omnibias.jax.last_linear_block
• omnibias.torch.last_linear_block
• omnibias.core.block_search.last_linear_block
• omnibias.jax.optim.last_linear_block
• omnibias.jax.optim_block_search.last_linear_block
• omnibias.torch.optim.last_linear_block
• omnibias.torch.optim_block_search.last_linear_block

FUNCTION

ombu_bias_block

omnibias.core.ombu_bias_blockSource ↗

ombu_bias_block(*, n_channels: int, k_biases: int, channel: int=0, slot: int=0, offset: int=0) -> BlockSpec

One OMBU bias entry in a row-major (n_channels, K) bank.

Parameters

Name	Type & default	Contract
n_channelskeyword-only	intRequired	See the operation contract and implementation below.
k_biaseskeyword-only	intRequired	See the operation contract and implementation below.
channelkeyword-only	intDefault: 0	See the operation contract and implementation below.
slotkeyword-only	intDefault: 0	See the operation contract and implementation below.
offsetkeyword-only	intDefault: 0	See the operation contract and implementation below.

ReturnsBlockSpec

Executed regression · ombu bias is one slot

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_ombu_bias_is_one_slot() -> None:
    spec = ombu_bias_block(n_channels=2, k_biases=3, channel=1, slot=1)
    flags = resolve_block_mask(6, spec)
    assert flags == (False, False, False, False, True, False)

python -m pytest "packages/omnibias-core/tests/test_block_search.py::test_ombu_bias_is_one_slot" -q

Regression example · sweep and named ombu mask

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

Read the implementation

def ombu_bias_block(
    *,
    n_channels: int,
    k_biases: int,
    channel: int = 0,
    slot: int = 0,
    offset: int = 0,
) -> BlockSpec:
    """One OMBU bias entry in a row-major ``(n_channels, K)`` bank."""
    return BlockSpec(
        kind="ombu_bias",
        n_channels=int(n_channels),
        k_biases=int(k_biases),
        index=int(channel),
        slot=int(slot),
        offset=int(offset),
    )

Related APIsomnibias.core.BlockSpec↗

Other public import paths

• omnibias.core.ombu_bias_block
• omnibias.jax.ombu_bias_block
• omnibias.torch.ombu_bias_block
• omnibias.core.block_search.ombu_bias_block
• omnibias.jax.optim.ombu_bias_block
• omnibias.jax.optim_block_search.ombu_bias_block
• omnibias.torch.optim.ombu_bias_block
• omnibias.torch.optim_block_search.ombu_bias_block

FUNCTION

resolve_block_mask

omnibias.core.resolve_block_maskSource ↗

resolve_block_mask(n_params: int, spec: BlockSpec, mask: Sequence[bool] | None=None) -> tuple[bool, ...]

Boolean mask of length n_params for spec.

Parameters

Name	Type & default	Contract
n_paramspositional-or-keyword	intRequired	See the operation contract and implementation below.
specpositional-or-keyword	BlockSpecRequired	See the operation contract and implementation below.
maskpositional-or-keyword	Sequence[bool] | NoneDefault: None	See the operation contract and implementation below.

Returnstuple[bool, ...]

Explicit validation & failure cases

• When n < 1, raises ValueError(f'n_params must be >= 1, got {n_params}').
• When mask is None, raises ValueError("kind='mask' requires a boolean mask").
• When len(mask) != n, raises ValueError(f'mask length {len(mask)} != n_params {n}').
• When start < 0 or stop > n or start >= stop, raises ValueError(f'last_linear slice [{start}, {stop}) is outside 0..{n}').
• When loc >= n, raises ValueError(f'ombu_bias index {loc} >= n_params {n}').
• When stop > n, raises ValueError(f'arrangement_W slice [{start}, {stop}) exceeds n_params {n}').

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

Executed regression · last linear mask is the tail

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_last_linear_mask_is_the_tail() -> None:
    spec = last_linear_block(2)
    assert resolve_block_mask(5, spec) == (False, False, False, True, True)

python -m pytest "packages/omnibias-core/tests/test_block_search.py::test_last_linear_mask_is_the_tail" -q

Regression example · arrangement w is one row

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_arrangement_w_is_one_row() -> None:
    spec = arrangement_w_block(n_hyperplanes=2, n_features=3, row=1)
    flags = resolve_block_mask(6, spec)
    assert flags == (False, False, False, True, True, True)

python -m pytest packages/omnibias-core/tests/test_block_search.py -k test_arrangement_w_is_one_row -q

Read the implementation

def resolve_block_mask(
    n_params: int,
    spec: BlockSpec,
    mask: Sequence[bool] | None = None,
) -> tuple[bool, ...]:
    """Boolean mask of length ``n_params`` for ``spec``."""
    n = int(n_params)
    if n < 1:
        raise ValueError(f"n_params must be >= 1, got {n_params}")
    flags = [False] * n
    if spec.kind == "mask":
        if mask is None:
            raise ValueError("kind='mask' requires a boolean mask")
        if len(mask) != n:
            raise ValueError(f"mask length {len(mask)} != n_params {n}")
        return tuple(bool(v) for v in mask)
    if spec.kind == "last_linear":
        width = int(spec.width or 0)
        start = n - width - int(spec.offset)
        stop = n - int(spec.offset)
        if start < 0 or stop > n or start >= stop:
            raise ValueError(
                f"last_linear slice [{start}, {stop}) is outside 0..{n}"
            )
        for i in range(start, stop):
            flags[i] = True
        return tuple(flags)
    if spec.kind == "ombu_bias":
        loc = int(spec.offset) + int(spec.index) * int(spec.k_biases or 0) + int(spec.slot)
        if loc >= n:
            raise ValueError(f"ombu_bias index {loc} >= n_params {n}")
        flags[loc] = True
        return tuple(flags)
    width = int(spec.n_features or 0)
    start = int(spec.offset) + int(spec.index) * width
    stop = start + width
    if stop > n:
        raise ValueError(f"arrangement_W slice [{start}, {stop}) exceeds n_params {n}")
    for i in range(start, stop):
        flags[i] = True
    return tuple(flags)

Other public import paths

• omnibias.core.resolve_block_mask
• omnibias.core.block_search.resolve_block_mask

FUNCTION

unit_direction_from_mask

omnibias.core.unit_direction_from_maskSource ↗

unit_direction_from_mask(mask: Sequence[bool], probe: Sequence[float] | None=None) -> tuple[float, ...]

Unit vector supported on mask.

If probe is given (typically -grad), it is restricted to the mask and normalised. A zero restriction falls back to the first active standard-basis vector.

Parameters

Name	Type & default	Contract
maskpositional-or-keyword	Sequence[bool]Required	See the operation contract and implementation below.
probepositional-or-keyword	Sequence[float] | NoneDefault: None	See the operation contract and implementation below.

Returnstuple[float, ...]

Explicit validation & failure cases

• When not active, raises ValueError('block mask is empty').
• When len(probe) != len(mask), raises ValueError(f'probe length {len(probe)} != mask length {len(mask)}').

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 · empty mask rejected

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_rejected() -> None:
    with pytest.raises(ValueError, match="empty"):
        unit_direction_from_mask((False, False))
    with pytest.raises(ValueError):
        BlockSpec(kind="last_linear", width=0)

python -m pytest packages/omnibias-core/tests/test_block_search.py -k test_empty_mask_rejected -q

Read the implementation

def unit_direction_from_mask(
    mask: Sequence[bool],
    probe: Sequence[float] | None = None,
) -> tuple[float, ...]:
    """Unit vector supported on ``mask``.

    If ``probe`` is given (typically ``-grad``), it is restricted to the
    mask and normalised. A zero restriction falls back to the first
    active standard-basis vector.
    """
    active = [i for i, flag in enumerate(mask) if flag]
    if not active:
        raise ValueError("block mask is empty")
    direction = [0.0] * len(mask)
    if probe is not None:
        if len(probe) != len(mask):
            raise ValueError(f"probe length {len(probe)} != mask length {len(mask)}")
        restricted = [float(probe[i]) if mask[i] else 0.0 for i in range(len(mask))]
        norm = math.sqrt(sum(v * v for v in restricted))
        if norm > 0.0 and math.isfinite(norm):
            return tuple(v / norm for v in restricted)
    direction[active[0]] = 1.0
    return tuple(direction)

Other public import paths

• omnibias.core.unit_direction_from_mask
• omnibias.core.block_search.unit_direction_from_mask

← Package overviewRunnable learning examples →