omnibias.torch.stencil

0.5.0rc1Release candidatePython >=3.10Apache-2.0

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

import omnibias.torch.stencil

Module contract & mathematical conventions

Bias-collapse stencils.

The K-bias multi-bias unit f_K(z) = sum_k s_k * sigma(z + b_k) becomes the (K-1)-th derivative of sigma in a specific limit. The exact limit depends on the choice of stencil for b_k and the matching rescaling of the signs s_k. Two stencils are supported.

**Forward-difference stencil.**

b_k = b + (k - 1) * delta for k = 1 .. K s_k = (-1)^(K-k) * binom(K-1, k-1) / delta^(K-1)

Then f_K(z) -> sigma^(K-1)(z + b) as delta -> 0+ with first-order accuracy O(delta) for sigma in C^K.

**Central-difference stencil.**

b_k = b + (k - (K+1)/2) * delta for k = 1 .. K s_k = (-1)^(K-k) * binom(K-1, k-1) / delta^(K-1)

The same scaling, but biases are centred around b instead of starting there. Order improves to O(delta^2) for sigma in C^(K+1).

We expose plain torch.Tensor builders for the stencils so the unit can pre-compute them once per (K, stencil) combination and only rescale when delta is updated.

Terminology: this delta -> 0 limit is *the* **founding bias collapse** (K biases collapse onto one value, yielding sigma^(K-1)). Do not confuse it with **temperature collapse**, the beta -> inf penalty in omnibias.convex (a feasibility step, not a derivative). See docs/derivatives.md.

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

central_bias_offsets

omnibias.torch.stencil.central_bias_offsetsSource ↗

central_bias_offsets(K: int, delta: float, dtype: torch.dtype | None=None) -> Tensor

Symmetric offsets ((k - (K+1)/2) * delta)_k=1..K.

The mean is zero so the symmetric stencil is centred on b.

dtype defaults to torch.get_default_dtype().

Parameters

Name	Type & default	Contract
`K`positional-or-keyword	`int`Required	See the operation contract and implementation below.
`delta`positional-or-keyword	`float`Required	See the operation contract and implementation below.
`dtype`positional-or-keyword	`torch.dtype \| None`Default: None	See the operation contract and implementation below.

ReturnsTensor

Executed regression · central bias offsets follows default dtype

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_central_bias_offsets_follows_default_dtype(torch_default_float64) -> None:
    b = central_bias_offsets(K=3, delta=0.5)
    assert b.dtype == torch.float64

python -m pytest "packages/omnibias-torch/tests/test_audit_regressions.py::test_central_bias_offsets_follows_default_dtype" -q

Regression example · literal matches analytic in well conditioned regime

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_literal_matches_analytic_in_well_conditioned_regime(
    name: str, delta: float, K: int
) -> None:
    """For modest ``K`` and modest ``delta``, the literal forward in
    float64 still agrees with the closed-form derivative to the
    truncation accuracy of the central-difference stencil."""
    spec = get_activation(name)
    z = torch.linspace(-1.0, 1.0, 32).unsqueeze(-1).expand(32, 4).contiguous().double()

    biases = central_bias_offsets(K, delta).double().unsqueeze(0).expand(4, K).contiguous()
    signs = central_difference_signs(K, delta).double().unsqueeze(0).expand(4, K).contiguous()

    literal = multibias_literal_forward(z, biases, signs, spec.forward)
    analytic = spec.fastpath(z, K - 1)
    rel = (literal - analytic).abs().max().item() / (analytic.abs().max().item() + 1e-12)
    # Central-difference stencil has truncation error O(delta^2); for the
    # well-conditioned regime (K <= 3, delta >= 1e-2) cancellation is
    # negligible so the empirical agreement is set by the truncation.
    assert rel < 5e-2, f"{name} K={K} delta={delta}: rel={rel:.2e}"

python -m pytest packages/omnibias-torch/tests/test_fastpath_stability.py -k test_literal_matches_analytic_in_well_conditioned_regime -q

Read the implementation

def central_bias_offsets(K: int, delta: float, dtype: torch.dtype | None = None) -> Tensor:
    """Symmetric offsets ``((k - (K+1)/2) * delta)_k=1..K``.

    The mean is zero so the symmetric stencil is centred on ``b``.

    ``dtype`` defaults to :func:`torch.get_default_dtype()`.
    """
    if dtype is None:
        dtype = torch.get_default_dtype()
    centre = (K + 1) / 2.0
    return (torch.arange(1, K + 1, dtype=dtype) - centre) * delta

FUNCTION

central_difference_signs

omnibias.torch.stencil.central_difference_signsSource ↗

central_difference_signs(K: int, delta: float, dtype: torch.dtype | None=None) -> Tensor

Central-difference signs.

Same magnitudes as forward_difference_signs; the difference is in where the biases are placed (see central_bias_offsets).

Parameters

Name	Type & default	Contract
`K`positional-or-keyword	`int`Required	See the operation contract and implementation below.
`delta`positional-or-keyword	`float`Required	See the operation contract and implementation below.
`dtype`positional-or-keyword	`torch.dtype \| None`Default: None	See the operation contract and implementation below.

ReturnsTensor

Executed regression · fastpath matches literal in collapse stencil

Full file, imports & fixtures ↗

def test_fastpath_matches_literal_in_collapse_stencil(name: str, Ks: list[int]) -> None:
    """At delta = 1e-2 the literal stencil forward agrees with the
    closed-form ``sigma^(K-1)`` to ``rtol=1e-4`` (limited by the O(delta^2)
    stencil error, not by the fast-path itself).
    """
    spec = get_activation(name)
    delta = 1e-2
    z = torch.linspace(-1.0, 1.0, 32).unsqueeze(-1).expand(32, 4).contiguous().double()
    bias_value = 0.0
    for K in Ks:
        biases, signs = _stencil_ombu_args(4, K, delta, bias_value=bias_value)
        biases = biases.double()
        signs = signs.double()
        literal = multibias_literal_forward(z, biases, signs, spec.forward)
        analytic = spec.fastpath(z + bias_value, K - 1)
        rel = (literal - analytic).abs().max().item() / (analytic.abs().max().item() + 1e-12)
        # The central-difference stencil at delta=1e-2 has truncation error
        # O(delta^2) = 1e-4 for smooth activations.
        assert rel < 5e-4, f"{name} K={K}: rel={rel:.2e}"

python -m pytest "packages/omnibias-torch/tests/test_fastpath_correctness.py::test_fastpath_matches_literal_in_collapse_stencil[exp-Ks4]" -q

Regression example · literal matches analytic in well conditioned regime

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_literal_matches_analytic_in_well_conditioned_regime(
    name: str, delta: float, K: int
) -> None:
    """For modest ``K`` and modest ``delta``, the literal forward in
    float64 still agrees with the closed-form derivative to the
    truncation accuracy of the central-difference stencil."""
    spec = get_activation(name)
    z = torch.linspace(-1.0, 1.0, 32).unsqueeze(-1).expand(32, 4).contiguous().double()

    biases = central_bias_offsets(K, delta).double().unsqueeze(0).expand(4, K).contiguous()
    signs = central_difference_signs(K, delta).double().unsqueeze(0).expand(4, K).contiguous()

    literal = multibias_literal_forward(z, biases, signs, spec.forward)
    analytic = spec.fastpath(z, K - 1)
    rel = (literal - analytic).abs().max().item() / (analytic.abs().max().item() + 1e-12)
    # Central-difference stencil has truncation error O(delta^2); for the
    # well-conditioned regime (K <= 3, delta >= 1e-2) cancellation is
    # negligible so the empirical agreement is set by the truncation.
    assert rel < 5e-2, f"{name} K={K} delta={delta}: rel={rel:.2e}"

python -m pytest packages/omnibias-torch/tests/test_fastpath_stability.py -k test_literal_matches_analytic_in_well_conditioned_regime -q

Read the implementation

def central_difference_signs(
    K: int, delta: float, dtype: torch.dtype | None = None
) -> Tensor:
    """Central-difference signs.

    Same magnitudes as :func:`forward_difference_signs`; the difference is
    in where the biases are placed (see :func:`central_bias_offsets`).
    """
    return forward_difference_signs(K, delta, dtype=dtype)

FUNCTION

forward_bias_offsets

omnibias.torch.stencil.forward_bias_offsetsSource ↗

forward_bias_offsets(K: int, delta: float, dtype: torch.dtype | None=None) -> Tensor

Offsets (0, delta, 2*delta, ..., (K-1)*delta).

dtype defaults to torch.get_default_dtype().

Parameters

Name	Type & default	Contract
`K`positional-or-keyword	`int`Required	See the operation contract and implementation below.
`delta`positional-or-keyword	`float`Required	See the operation contract and implementation below.
`dtype`positional-or-keyword	`torch.dtype \| None`Default: None	See the operation contract and implementation below.

ReturnsTensor

Executed regression · forward bias offsets follows default dtype

Full file, imports & fixtures ↗

def test_forward_bias_offsets_follows_default_dtype(torch_default_float64) -> None:
    b = forward_bias_offsets(K=3, delta=0.5)
    assert b.dtype == torch.float64

python -m pytest "packages/omnibias-torch/tests/test_audit_regressions.py::test_forward_bias_offsets_follows_default_dtype" -q

Read the implementation

def forward_bias_offsets(K: int, delta: float, dtype: torch.dtype | None = None) -> Tensor:
    """Offsets ``(0, delta, 2*delta, ..., (K-1)*delta)``.

    ``dtype`` defaults to :func:`torch.get_default_dtype()`.
    """
    if dtype is None:
        dtype = torch.get_default_dtype()
    return torch.arange(K, dtype=dtype) * delta

FUNCTION

forward_difference_signs

omnibias.torch.stencil.forward_difference_signsSource ↗

forward_difference_signs(K: int, delta: float, dtype: torch.dtype | None=None) -> Tensor

Forward-difference signs at order K-1 with step delta.

Produces s_k = (-1)^(K-k) * binom(K-1, k-1) / delta^(K-1) so that sum_k s_k * sigma(z + b + (k-1)*delta) -> sigma^(K-1)(z + b) as delta -> 0+.

For K = 1 returns [1.0] (the standard single-bias unit).

dtype defaults to torch.get_default_dtype().

Parameters

Name	Type & default	Contract
`K`positional-or-keyword	`int`Required	See the operation contract and implementation below.
`delta`positional-or-keyword	`float`Required	See the operation contract and implementation below.
`dtype`positional-or-keyword	`torch.dtype \| None`Default: None	See the operation contract and implementation below.

ReturnsTensor

Explicit validation & failure cases

When K < 1, raises ValueError(f'K must be >= 1, got {K}').
When delta <= 0 and K > 1, raises ValueError(f'delta must be > 0 for K > 1 (got delta={delta}, K={K})').

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

Executed regression · forward difference signs follows default dtype

Full file, imports & fixtures ↗

def test_forward_difference_signs_follows_default_dtype(
    torch_default_float64, K
) -> None:
    s = forward_difference_signs(K, delta=0.1)
    assert s.dtype == torch.float64

python -m pytest "packages/omnibias-torch/tests/test_audit_regressions.py::test_forward_difference_signs_follows_default_dtype[1]" -q

Read the implementation

def forward_difference_signs(
    K: int, delta: float, dtype: torch.dtype | None = None
) -> Tensor:
    """Forward-difference signs at order K-1 with step ``delta``.

    Produces ``s_k = (-1)^(K-k) * binom(K-1, k-1) / delta^(K-1)`` so that
    ``sum_k s_k * sigma(z + b + (k-1)*delta) -> sigma^(K-1)(z + b)``
    as ``delta -> 0+``.

    For ``K = 1`` returns ``[1.0]`` (the standard single-bias unit).

    ``dtype`` defaults to :func:`torch.get_default_dtype()`.
    """
    if K < 1:
        raise ValueError(f"K must be >= 1, got {K}")
    if delta <= 0 and K > 1:
        raise ValueError(f"delta must be > 0 for K > 1 (got delta={delta}, K={K})")
    if dtype is None:
        dtype = torch.get_default_dtype()
    if K == 1:
        return torch.tensor([1.0], dtype=dtype)
    inv_scale = 1.0 / (delta ** (K - 1))
    signs = [((-1.0) ** (K - 1 - j)) * math.comb(K - 1, j) * inv_scale for j in range(K)]
    return torch.tensor(signs, dtype=dtype)

FUNCTION

identity_signs

omnibias.torch.identity_init.identity_signsSource ↗

identity_signs(K: int, dtype: torch.dtype | None=None) -> Tensor

Signs that satisfy Lemma 1 (identity nesting).

Returns s with sum(s) = 1.

For odd K, the canonical alternating pattern (+1, -1, +1, ..., +1) already sums to one.
For even K, alternating gives sum = 0; we double the first sign (so s[0] = +2) to bring the sum to one. The last K-1 entries remain strictly alternating, which keeps the unit balanced and matches the empirical pattern reported in the multi-bias paper's XOR construction (Figure 1, trained s ~ (+2.18, -4.42)).

dtype defaults to torch.get_default_dtype() so callers that have set a non-default global dtype (e.g. float64 for scientific work) get matching parameters by default.

Parameters

Name	Type & default	Contract
`K`positional-or-keyword	`int`Required	See the operation contract and implementation below.
`dtype`positional-or-keyword	`torch.dtype \| None`Default: None	See the operation contract and implementation below.

ReturnsTensor

Explicit validation & failure cases

When K < 1, raises ValueError(f'K must be >= 1, got {K}').

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

Executed regression · identity signs follows default dtype

Full file, imports & fixtures ↗

def test_identity_signs_follows_default_dtype(torch_default_float64) -> None:
    s = identity_signs(K=3)
    assert s.dtype == torch.float64

python -m pytest "packages/omnibias-torch/tests/test_audit_regressions.py::test_identity_signs_follows_default_dtype" -q

Regression example · dtype default passthrough does not break float32

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_dtype_default_passthrough_does_not_break_float32(monkeypatch) -> None:
    """When default dtype is float32 (the global default), helpers still
    produce float32 (no regression)."""
    assert torch.get_default_dtype() == torch.float32
    biases = identity_init_biases(num_channels=4, K=2)
    assert biases.dtype == torch.float32
    s = identity_signs(K=3)
    assert s.dtype == torch.float32

python -m pytest packages/omnibias-torch/tests/test_audit_regressions.py -k test_dtype_default_passthrough_does_not_break_float32 -q

Read the implementation

def identity_signs(K: int, dtype: torch.dtype | None = None) -> Tensor:
    """Signs that satisfy Lemma 1 (identity nesting).

    Returns ``s`` with ``sum(s) = 1``.

    - For odd K, the canonical alternating pattern ``(+1, -1, +1, ..., +1)``
      already sums to one.
    - For even K, alternating gives ``sum = 0``; we double the first sign
      (so ``s[0] = +2``) to bring the sum to one. The last K-1 entries
      remain strictly alternating, which keeps the unit balanced and
      matches the empirical pattern reported in the multi-bias paper's
      XOR construction (Figure 1, trained ``s ~ (+2.18, -4.42)``).

    ``dtype`` defaults to :func:`torch.get_default_dtype()` so callers
    that have set a non-default global dtype (e.g. ``float64`` for
    scientific work) get matching parameters by default.
    """
    if K < 1:
        raise ValueError(f"K must be >= 1, got {K}")
    if dtype is None:
        dtype = torch.get_default_dtype()
    s = torch.tensor([(-1.0) ** k for k in range(K)], dtype=dtype)  # +1, -1, +1, ...
    if K % 2 == 0:
        s[0] = 2.0  # makes sum = 1 while preserving alternation in s[1:]
    return s

Other public import paths

omnibias.torch.identity_init.identity_signs
omnibias.torch.stencil.identity_signs

FUNCTION

stencil_offsets

omnibias.torch.stencil.stencil_offsetsSource ↗

stencil_offsets(K: int, delta: float, stencil: Literal['forward', 'central']='central', dtype: torch.dtype | None=None) -> Tensor

Bias offsets for the chosen stencil.

Parameters

Name	Type & default	Contract
`K`positional-or-keyword	`int`Required	See the operation contract and implementation below.
`delta`positional-or-keyword	`float`Required	See the operation contract and implementation below.
`stencil`positional-or-keyword	`Literal['forward', 'central']`Default: 'central'	See the operation contract and implementation below.
`dtype`positional-or-keyword	`torch.dtype \| None`Default: None	See the operation contract and implementation below.

ReturnsTensor

Executed regression · stencil offsets follows default dtype

Full file, imports & fixtures ↗

def test_stencil_offsets_follows_default_dtype(torch_default_float64) -> None:
    for stencil in ("forward", "central"):
        b = stencil_offsets(K=3, delta=0.5, stencil=stencil)
        assert b.dtype == torch.float64, stencil

python -m pytest "packages/omnibias-torch/tests/test_audit_regressions.py::test_stencil_offsets_follows_default_dtype" -q

Regression example · forward bias offsets follows default dtype

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_forward_bias_offsets_follows_default_dtype(torch_default_float64) -> None:
    b = forward_bias_offsets(K=3, delta=0.5)
    assert b.dtype == torch.float64

python -m pytest packages/omnibias-torch/tests/test_audit_regressions.py -k test_forward_bias_offsets_follows_default_dtype -q

Read the implementation

def stencil_offsets(
    K: int,
    delta: float,
    stencil: Literal["forward", "central"] = "central",
    dtype: torch.dtype | None = None,
) -> Tensor:
    """Bias offsets for the chosen stencil."""
    if stencil == "forward":
        return forward_bias_offsets(K, delta, dtype=dtype)
    if stencil == "central":
        return central_bias_offsets(K, delta, dtype=dtype)
    raise ValueError(f"Unknown stencil {stencil!r}; expected 'forward' or 'central'.")

FUNCTION

stencil_signs

omnibias.torch.stencil.stencil_signsSource ↗

stencil_signs(K: int, delta: float, stencil: Literal['forward', 'central']='central', dtype: torch.dtype | None=None) -> Tensor

Signs matched to the chosen stencil (same magnitudes for both).

Parameters

Name	Type & default	Contract
`K`positional-or-keyword	`int`Required	See the operation contract and implementation below.
`delta`positional-or-keyword	`float`Required	See the operation contract and implementation below.
`stencil`positional-or-keyword	`Literal['forward', 'central']`Default: 'central'	See the operation contract and implementation below.
`dtype`positional-or-keyword	`torch.dtype \| None`Default: None	See the operation contract and implementation below.

ReturnsTensor

Explicit validation & failure cases

When stencil not in ('forward', 'central'), raises ValueError(f"Unknown stencil {stencil!r}; expected 'forward' or 'central'.").

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

Executed regression · stencil signs follows default dtype

Full file, imports & fixtures ↗

def test_stencil_signs_follows_default_dtype(torch_default_float64) -> None:
    for stencil in ("forward", "central"):
        s = stencil_signs(K=3, delta=0.5, stencil=stencil)
        assert s.dtype == torch.float64, stencil

python -m pytest "packages/omnibias-torch/tests/test_audit_regressions.py::test_stencil_signs_follows_default_dtype" -q

Regression example · identity signs follows default dtype

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_identity_signs_follows_default_dtype(torch_default_float64) -> None:
    s = identity_signs(K=3)
    assert s.dtype == torch.float64

python -m pytest packages/omnibias-torch/tests/test_audit_regressions.py -k test_identity_signs_follows_default_dtype -q

Read the implementation

def stencil_signs(
    K: int,
    delta: float,
    stencil: Literal["forward", "central"] = "central",
    dtype: torch.dtype | None = None,
) -> Tensor:
    """Signs matched to the chosen stencil (same magnitudes for both)."""
    if stencil not in ("forward", "central"):
        raise ValueError(f"Unknown stencil {stencil!r}; expected 'forward' or 'central'.")
    return forward_difference_signs(K, delta, dtype=dtype)

← Package overview Runnable learning examples →

omnibias.torch.stencil

central_bias_offsets​

Parameters

central_difference_signs​

Parameters

forward_bias_offsets​

Parameters

forward_difference_signs​

Parameters

identity_signs​

Parameters

stencil_offsets​

Parameters

stencil_signs​

Parameters

central_bias_offsets

central_difference_signs

forward_bias_offsets

forward_difference_signs

identity_signs

stencil_offsets

stencil_signs