API REFERENCE/omnibias-torch

omnibias.torch.fastpath.dispatch

0.5.0rc1Release candidatePython >=3.10Apache-2.0

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

import omnibias.torch.fastpath.dispatch

Module contract & mathematical conventions

Forward-pass helpers used by OperatorMultiBiasUnit and OperatorBlock.

Two kernels:

• multibias_literal_forward: sum_k s_k * sigma(z + b_k), the literal multi-bias activation.
• multibias_collapsed_forward: closed-form sigma^(order)(z + bias_mean) via the activation spec's fast-path kernel; one base-activation evaluation regardless of K.
• multibias_integral_forward: closed-form sign-weighted antiderivative sum, whose K=2 alternating-sign form is a definite integral over a bias interval.

Both are plain functions (no nn.Module overhead) so they can be called from inside training loops without parameter bookkeeping.

**Why no custom autograd Function?** The bias-gradient scrambling that the predecessor multibias v0.2.1 release patched was a consequence of computing (sigma(z + b + delta) - sigma(z + b)) / delta in float32 and then asking autograd to differentiate that quotient: the forward already lost log10(1/delta) digits, and the backward lost the rest. The closed-form fast paths in omnibias.fastpath avoid the quotient entirely -- sigma^(n)(z + bar b) is one sigmoid plus a Horner polynomial in the result -- and autograd through Horner is bit-stable. There is therefore nothing for a custom backward to fix in the collapsed branch; we let autograd do its job.

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 →

FUNCTION

bias_spread

omnibias.torch.fastpath.bias_spreadSource ↗

bias_spread(biases: Tensor) -> Tensor

Per-channel max(b_k) - min(b_k); near zero in the collapse regime.

Parameters

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

ReturnsTensor

Read the implementation

def bias_spread(biases: Tensor) -> Tensor:
    """Per-channel ``max(b_k) - min(b_k)``; near zero in the collapse regime."""
    return biases.max(dim=-1).values - biases.min(dim=-1).values

Other public import paths

• omnibias.torch.fastpath.bias_spread
• omnibias.torch.fastpath.dispatch.bias_spread

FUNCTION

is_collapsed

omnibias.torch.fastpath.is_collapsedSource ↗

is_collapsed(biases: Tensor, threshold: float) -> bool

True iff every channel's bias spread is strictly below threshold.

Parameters

Name	Type & default	Contract
biasespositional-or-keyword	TensorRequired	See the operation contract and implementation below.
thresholdpositional-or-keyword	floatRequired	See the operation contract and implementation below.

Returnsbool

Read the implementation

def is_collapsed(biases: Tensor, threshold: float) -> bool:
    """True iff every channel's bias spread is strictly below ``threshold``."""
    if threshold <= 0:
        return bool((bias_spread(biases) == 0).all().item())
    return bool((bias_spread(biases) < threshold).all().item())

Related APIsomnibias.torch.fastpath.bias_spread↗

Other public import paths

• omnibias.torch.fastpath.is_collapsed
• omnibias.torch.fastpath.dispatch.is_collapsed

FUNCTION

multibias_collapsed_forward

omnibias.torch.fastpath.multibias_collapsed_forwardSource ↗

multibias_collapsed_forward(z: Tensor, biases: Tensor, spec: ActivationSpec[Tensor], order: int) -> Tensor

Closed-form sigma^(order)(z + bias_mean).

Parameters

z : Tensor of shape (..., C) biases : Tensor of shape (C, K) Per-channel biases; only their per-channel mean is used. spec : ActivationSpec Must have a non-None fastpath kernel. order : int Derivative order n; must be >= 0.

Returns

Tensor of shape (..., C).

Parameters

Name	Type & default	Contract
zpositional-or-keyword	TensorRequired	Tensor of shape (..., C)
biasespositional-or-keyword	TensorRequired	Tensor of shape (C, K) Per-channel biases; only their per-channel mean is used.
specpositional-or-keyword	ActivationSpec[Tensor]Required	ActivationSpec Must have a non-None fastpath kernel.
orderpositional-or-keyword	intRequired	int Derivative order n; must be >= 0.

ReturnsTensor

Explicit validation & failure cases

• When order < 0, raises ValueError(f'order must be >= 0, got {order}.').
• When spec.fastpath is None, raises NotImplementedError(f'Activation {spec.name!r} has no closed-form derivative kernel.').

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

Read the implementation

def multibias_collapsed_forward(
    z: Tensor,
    biases: Tensor,
    spec: ActivationSpec[Tensor],
    order: int,
) -> Tensor:
    """Closed-form ``sigma^(order)(z + bias_mean)``.

    Parameters
    ----------
    z : Tensor of shape ``(..., C)``
    biases : Tensor of shape ``(C, K)``
        Per-channel biases; only their per-channel mean is used.
    spec : :class:`ActivationSpec`
        Must have a non-None ``fastpath`` kernel.
    order : int
        Derivative order ``n``; must be ``>= 0``.

    Returns
    -------
    Tensor of shape ``(..., C)``.
    """
    if order < 0:
        raise ValueError(f"order must be >= 0, got {order}.")
    if spec.fastpath is None:
        raise NotImplementedError(f"Activation {spec.name!r} has no closed-form derivative kernel.")
    b_mean = biases.mean(dim=-1)  # (C,)
    out: Tensor = spec.fastpath(z + b_mean, order)
    return out

Other public import paths

• omnibias.torch.fastpath.multibias_collapsed_forward
• omnibias.torch.fastpath.dispatch.multibias_collapsed_forward

FUNCTION

multibias_forward

omnibias.torch.fastpath.multibias_forwardSource ↗

multibias_forward(z: Tensor, biases: Tensor, signs: Tensor, spec: ActivationSpec[Tensor], *, collapse_threshold: float=0.0, use_collapsed: bool=False, order: int | None=None) -> Tensor

Combined dispatcher.

If use_collapsed is True (and the spec has a fast-path kernel), computes multibias_collapsed_forward at the requested order (default biases.shape[-1] - 1). Otherwise computes the literal forward.

.. deprecated:: collapse_threshold never affected dispatch and is deprecated. It is kept only for call-site compatibility; passing a non-zero value now emits a DeprecationWarning. To select the collapsed fast-path based on the bias spread, test is_collapsed yourself and pass use_collapsed=True explicitly.

Parameters

Name	Type & default	Contract
zpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
biasespositional-or-keyword	TensorRequired	See the operation contract and implementation below.
signspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
specpositional-or-keyword	ActivationSpec[Tensor]Required	See the operation contract and implementation below.
collapse_thresholdkeyword-only	floatDefault: 0.0	See the operation contract and implementation below.
use_collapsedkeyword-only	boolDefault: False	See the operation contract and implementation below.
orderkeyword-only	int | NoneDefault: None	See the operation contract and implementation below.

ReturnsTensor

Executed regression · multibias forward collapse threshold is deprecated noop

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_multibias_forward_collapse_threshold_is_deprecated_noop() -> None:
    """``collapse_threshold`` never affected dispatch: a non-zero value warns
    (DeprecationWarning) but leaves the result equal to the literal forward,
    and the default value is silent."""
    spec = get_activation("sigmoid")
    z = torch.linspace(-1.0, 1.0, 6).unsqueeze(-1).double()  # (6, 1)
    biases = torch.zeros(1, 3, dtype=torch.float64)
    signs = torch.tensor([[1.0, -1.0, 1.0]], dtype=torch.float64)
    literal = multibias_literal_forward(z, biases, signs, spec.forward)

    with pytest.warns(DeprecationWarning, match="collapse_threshold"):
        out = multibias_forward(z, biases, signs, spec, collapse_threshold=1e-3)
    assert torch.equal(out, literal)

    with warnings.catch_warnings():
        warnings.simplefilter("error")  # any warning here would fail the test
        out_default = multibias_forward(z, biases, signs, spec)
    assert torch.equal(out_default, literal)

python -m pytest "packages/omnibias-torch/tests/test_fastpath_correctness.py::test_multibias_forward_collapse_threshold_is_deprecated_noop" -q

Read the implementation

def multibias_forward(
    z: Tensor,
    biases: Tensor,
    signs: Tensor,
    spec: ActivationSpec[Tensor],
    *,
    collapse_threshold: float = 0.0,
    use_collapsed: bool = False,
    order: int | None = None,
) -> Tensor:
    """Combined dispatcher.

    If ``use_collapsed`` is True (and the spec has a fast-path kernel),
    computes :func:`multibias_collapsed_forward` at the requested order
    (default ``biases.shape[-1] - 1``). Otherwise computes the literal
    forward.

    .. deprecated::
        ``collapse_threshold`` never affected dispatch and is deprecated. It
        is kept only for call-site compatibility; passing a non-zero value
        now emits a :class:`DeprecationWarning`. To select the collapsed
        fast-path based on the bias spread, test :func:`is_collapsed`
        yourself and pass ``use_collapsed=True`` explicitly.
    """
    if collapse_threshold != 0.0:
        warnings.warn(
            "multibias_forward(collapse_threshold=...) is a no-op and is "
            "deprecated: it has never influenced which forward is selected. "
            "Use is_collapsed(biases, threshold) to test the collapse regime "
            "and pass use_collapsed=True to choose the collapsed fast-path.",
            DeprecationWarning,
            stacklevel=2,
        )
    if use_collapsed and spec.fastpath is not None:
        K = biases.shape[-1]
        return multibias_collapsed_forward(z, biases, spec, K - 1 if order is None else order)
    return multibias_literal_forward(z, biases, signs, spec.forward)

Related APIsomnibias.torch.fastpath.multibias_collapsed_forward↗omnibias.torch.fastpath.multibias_literal_forward↗

Other public import paths

• omnibias.torch.fastpath.multibias_forward
• omnibias.torch.fastpath.dispatch.multibias_forward

FUNCTION

multibias_integral_forward

omnibias.torch.fastpath.dispatch.multibias_integral_forwardSource ↗

multibias_integral_forward(z: Tensor, biases: Tensor, signs: Tensor, spec: ActivationSpec[Tensor]) -> Tensor

Closed-form sign-weighted antiderivative forward.

Computes sum_k signs[c, k] * S(z[..., c] + biases[c, k]) where S' = sigma. For K=2 with signs (+1, -1) this is the oriented integral S(z + b_1) - S(z + b_2) across the bias window.

Parameters

Name	Type & default	Contract
zpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
biasespositional-or-keyword	TensorRequired	See the operation contract and implementation below.
signspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
specpositional-or-keyword	ActivationSpec[Tensor]Required	See the operation contract and implementation below.

ReturnsTensor

Explicit validation & failure cases

• When spec.integral is None, raises NotImplementedError(f'Activation {spec.name!r} has no closed-form integral kernel; use literal_forward() instead, or pick a base with an antiderivative.').

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

Executed regression · integral block zero delta is blank window

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_integral_block_zero_delta_is_blank_window() -> None:
    """A zero-width strict integral window initializes as a blank feature."""
    block = OperatorBlock(op="integral", base="sigmoid", channels=3, init_delta=0.0)
    z = torch.randn(2, 3)
    out = block(z)
    assert torch.allclose(out, torch.zeros_like(out), atol=1e-6)

python -m pytest "packages/omnibias-torch/tests/test_blocks.py::test_integral_block_zero_delta_is_blank_window" -q

Read the implementation

def multibias_integral_forward(
    z: Tensor,
    biases: Tensor,
    signs: Tensor,
    spec: ActivationSpec[Tensor],
) -> Tensor:
    """Closed-form sign-weighted antiderivative forward.

    Computes ``sum_k signs[c, k] * S(z[..., c] + biases[c, k])`` where
    ``S' = sigma``. For K=2 with signs ``(+1, -1)`` this is the oriented
    integral ``S(z + b_1) - S(z + b_2)`` across the bias window.
    """
    if spec.integral is None:
        raise NotImplementedError(
            f"Activation {spec.name!r} has no closed-form integral kernel; "
            "use literal_forward() instead, or pick a base with an antiderivative."
        )
    return multibias_literal_forward(z, biases, signs, spec.integral)

Related APIsomnibias.torch.fastpath.multibias_literal_forward↗

FUNCTION

multibias_integral_window_forward

omnibias.torch.fastpath.dispatch.multibias_integral_window_forwardSource ↗

multibias_integral_window_forward(z: Tensor, center: Tensor, width: Tensor, signs: Tensor, spec: ActivationSpec[Tensor], *, normalize: bool=False, small_width_threshold: float=0.0001, use_small_width_taylor: bool=True) -> Tensor

Closed-form ordered bias-window integral.

center and positive width define endpoints b_lo = center - width / 2 and b_hi = center + width / 2. With fixed signs (-1, +1) this computes S(z + b_hi) - S(z + b_lo), i.e. the low-to-high definite integral of sigma across the bias interval. normalize=True returns the window average instead of total area.

Parameters

Name	Type & default	Contract
zpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
centerpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
widthpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
signspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
specpositional-or-keyword	ActivationSpec[Tensor]Required	See the operation contract and implementation below.
normalizekeyword-only	boolDefault: False	See the operation contract and implementation below.
small_width_thresholdkeyword-only	floatDefault: 0.0001	See the operation contract and implementation below.
use_small_width_taylorkeyword-only	boolDefault: True	See the operation contract and implementation below.

ReturnsTensor

Explicit validation & failure cases

• When width.shape != center.shape, raises ValueError(f'width shape {tuple(width.shape)} must match center shape {tuple(center.shape)}.').
• When signs.shape != (center.shape[0], 2), raises ValueError(f'signs shape {tuple(signs.shape)} must be ({center.shape[0]}, 2).').

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

Executed regression · integral block zero delta is blank window

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_integral_block_zero_delta_is_blank_window() -> None:
    """A zero-width strict integral window initializes as a blank feature."""
    block = OperatorBlock(op="integral", base="sigmoid", channels=3, init_delta=0.0)
    z = torch.randn(2, 3)
    out = block(z)
    assert torch.allclose(out, torch.zeros_like(out), atol=1e-6)

python -m pytest "packages/omnibias-torch/tests/test_blocks.py::test_integral_block_zero_delta_is_blank_window" -q

Read the implementation

def multibias_integral_window_forward(
    z: Tensor,
    center: Tensor,
    width: Tensor,
    signs: Tensor,
    spec: ActivationSpec[Tensor],
    *,
    normalize: bool = False,
    small_width_threshold: float = 1e-4,
    use_small_width_taylor: bool = True,
) -> Tensor:
    """Closed-form ordered bias-window integral.

    ``center`` and positive ``width`` define endpoints
    ``b_lo = center - width / 2`` and ``b_hi = center + width / 2``. With
    fixed signs ``(-1, +1)`` this computes
    ``S(z + b_hi) - S(z + b_lo)``, i.e. the low-to-high definite integral
    of ``sigma`` across the bias interval. ``normalize=True`` returns the
    window average instead of total area.
    """
    if width.shape != center.shape:
        raise ValueError(f"width shape {tuple(width.shape)} must match center shape {tuple(center.shape)}.")
    if signs.dim() == 1:
        signs = signs.unsqueeze(0)
    if signs.shape != (center.shape[0], 2):
        raise ValueError(f"signs shape {tuple(signs.shape)} must be ({center.shape[0]}, 2).")

    half_width = 0.5 * width
    endpoints = torch.stack((center - half_width, center + half_width), dim=-1)
    exact = multibias_integral_forward(z, endpoints, signs, spec)
    if normalize:
        exact = exact / width.clamp_min(width.new_tensor(1e-12))

    if not use_small_width_taylor or small_width_threshold <= 0.0:
        return exact

    midpoint = z + center
    approx = spec.forward(midpoint)
    if not normalize:
        approx = approx * width
    mask = width < small_width_threshold
    return torch.where(mask, approx, exact)

Related APIsomnibias.torch.fastpath.dispatch.multibias_integral_forward↗

FUNCTION

multibias_literal_forward

omnibias.torch.fastpath.multibias_literal_forwardSource ↗

multibias_literal_forward(z: Tensor, biases: Tensor, signs: Tensor, sigma: Callable[[Tensor], Tensor]) -> Tensor

Literal multi-bias activation.

Computes out[..., c] = sum_k signs[c, k] * sigma(z[..., c] + biases[c, k]).

Parameters

z : Tensor of shape (..., C) Input pre-activations, channel-last. biases : Tensor of shape (C, K) Per-channel bias terms. signs : Tensor of shape (C, K) or (K,) Per-channel signs (or shared across channels). sigma : Callable Base activation, applied elementwise.

Returns

Tensor of shape (..., C).

Parameters

Name	Type & default	Contract
zpositional-or-keyword	TensorRequired	Tensor of shape (..., C) Input pre-activations, channel-last.
biasespositional-or-keyword	TensorRequired	Tensor of shape (C, K) Per-channel bias terms.
signspositional-or-keyword	TensorRequired	Tensor of shape (C, K) or (K,) Per-channel signs (or shared across channels).
sigmapositional-or-keyword	Callable[[Tensor], Tensor]Required	Callable Base activation, applied elementwise.

ReturnsTensor

Explicit validation & failure cases

• When z.shape[-1] != biases.shape[0], raises ValueError(f'Input last-dim ({z.shape[-1]}) must equal num_channels ({biases.shape[0]}).').
• When signs.shape != biases.shape, raises ValueError(f'signs shape {tuple(signs.shape)} does not match biases shape {tuple(biases.shape)}.').

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

Executed regression · cmblinear forward shape

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_cmblinear_forward_shape() -> None:
    layer = cmbLinear(5, 3, op="identity", base="tanh")
    out = layer(torch.randn(2, 5))
    assert out.shape == (2, 3)

python -m pytest "packages/omnibias-torch/tests/test_blocks.py::test_cmblinear_forward_shape" -q

Regression example · multibias forward collapse threshold is deprecated noop

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_multibias_forward_collapse_threshold_is_deprecated_noop() -> None:
    """``collapse_threshold`` never affected dispatch: a non-zero value warns
    (DeprecationWarning) but leaves the result equal to the literal forward,
    and the default value is silent."""
    spec = get_activation("sigmoid")
    z = torch.linspace(-1.0, 1.0, 6).unsqueeze(-1).double()  # (6, 1)
    biases = torch.zeros(1, 3, dtype=torch.float64)
    signs = torch.tensor([[1.0, -1.0, 1.0]], dtype=torch.float64)
    literal = multibias_literal_forward(z, biases, signs, spec.forward)

    with pytest.warns(DeprecationWarning, match="collapse_threshold"):
        out = multibias_forward(z, biases, signs, spec, collapse_threshold=1e-3)
    assert torch.equal(out, literal)

    with warnings.catch_warnings():
        warnings.simplefilter("error")  # any warning here would fail the test
        out_default = multibias_forward(z, biases, signs, spec)
    assert torch.equal(out_default, literal)

python -m pytest packages/omnibias-torch/tests/test_fastpath_correctness.py -k test_multibias_forward_collapse_threshold_is_deprecated_noop -q

Read the implementation

def multibias_literal_forward(
    z: Tensor,
    biases: Tensor,
    signs: Tensor,
    sigma: Callable[[Tensor], Tensor],
) -> Tensor:
    """Literal multi-bias activation.

    Computes ``out[..., c] = sum_k signs[c, k] * sigma(z[..., c] + biases[c, k])``.

    Parameters
    ----------
    z : Tensor of shape ``(..., C)``
        Input pre-activations, channel-last.
    biases : Tensor of shape ``(C, K)``
        Per-channel bias terms.
    signs : Tensor of shape ``(C, K)`` or ``(K,)``
        Per-channel signs (or shared across channels).
    sigma : Callable
        Base activation, applied elementwise.

    Returns
    -------
    Tensor of shape ``(..., C)``.
    """
    if z.shape[-1] != biases.shape[0]:
        raise ValueError(
            f"Input last-dim ({z.shape[-1]}) must equal num_channels ({biases.shape[0]})."
        )
    if signs.dim() == 1:
        signs = signs.unsqueeze(0)  # (1, K) for broadcasting over channels
    elif signs.dim() == 2:
        if signs.shape != biases.shape:
            raise ValueError(
                f"signs shape {tuple(signs.shape)} does not match biases shape "
                f"{tuple(biases.shape)}."
            )
    else:
        raise ValueError(f"signs must be 1D or 2D, got shape {tuple(signs.shape)}.")

    # z: (..., C) -> (..., C, 1); biases: (C, K). Broadcast adds to (..., C, K).
    z_expand = z.unsqueeze(-1)
    activated = sigma(z_expand + biases)
    out = (activated * signs).sum(dim=-1)
    return out

Other public import paths

• omnibias.torch.fastpath.multibias_literal_forward
• omnibias.torch.fastpath.dispatch.multibias_literal_forward

← Package overviewRunnable learning examples →