API REFERENCE/omnibias-torch

omnibias.torch.scan

0.5.0rc1Release candidatePython >=3.10Apache-2.0

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

import omnibias.torch.scan

Module contract & mathematical conventions

Bias scan: shared template on a bank of offsets (torch; theory 01-02).

Scan[T](z)_j = F_T(z + tau_j). Parameters of T are shared; only the offset varies. Equivariance is along the transverse coordinate and exact on the bank lattice as an *interior* shift (tanh' is not periodic, so the response is not a circular wrap).

The template collapse is delta -> 0. Soft-argmax gamma is a softmax sharpness; gamma -> inf would be temperature collapse, a different limit.

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
BankSpec	omnibias.core.BankSpec

CLASS

BiasScan

omnibias.torch.BiasScanSource ↗

class BiasScan(nn.Module)

Channel-wise bias scan.

Parameters

num_channels: Trailing channel count of z. bank: Offset / scale bank. template: ~omnibias.core.multipack.MultiPackSpec or an operator role name. op: 01-13 catalog alias for template. Pass only one of template or op; they name the same six-role family, not a seventh role. readout: response -> (..., C, M); pooled / argmax -> (..., C). gamma: Soft-argmax sharpness (not a collapse parameter).

Parameters

Name	Type & default	Contract
num_channelspositional-or-keyword	intRequired	Trailing channel count of z.
bankpositional-or-keyword	BankSpecRequired	Offset / scale bank.
templatekeyword-only	MultiPackSpec | OpName | str | NoneDefault: None	~omnibias.core.multipack.MultiPackSpec or an operator role name.
opkeyword-only	OpName | str | NoneDefault: None	01-13 catalog alias for template. Pass only one of template or op; they name the same six-role family, not a seventh role.
basekeyword-only	str | ActivationSpec[Tensor]Default: 'tanh'	See the operation contract and implementation below.
learnable_offsetskeyword-only	boolDefault: True	See the operation contract and implementation below.
learnable_scaleskeyword-only	boolDefault: False	See the operation contract and implementation below.
readoutkeyword-only	ReadoutDefault: 'response'	response -> (..., C, M); pooled / argmax -> (..., C).
gammakeyword-only	floatDefault: 8.0	Soft-argmax sharpness (not a collapse parameter).
derivative_orderkeyword-only	intDefault: 1	See the operation contract and implementation below.
gapkeyword-only	floatDefault: 1.0	See the operation contract and implementation below.
dtypekeyword-only	torch.dtype | NoneDefault: None	See the operation contract and implementation below.

Executed regression · dtype none uses framework default

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_dtype_none_uses_framework_default() -> None:
    prev = torch.get_default_dtype()
    try:
        torch.set_default_dtype(torch.float64)
        scan = BiasScan(1, BankSpec.uniform(-1.0, 1.0, 3), dtype=None)
        assert scan.offsets.dtype == torch.float64
    finally:
        torch.set_default_dtype(prev)

python -m pytest "packages/omnibias-torch/tests/test_scan.py::test_dtype_none_uses_framework_default" -q

Regression example · min offset separation on learnable offsets

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_min_offset_separation_on_learnable_offsets() -> None:
    torch.set_default_dtype(torch.float64)
    bank = BankSpec.uniform(-1.0, 1.0, 5)
    scan = _scan(bank=bank, learnable_offsets=True)
    assert float(scan.min_offset_separation().detach()) == pytest.approx(0.5)
    with torch.no_grad():
        scan.offsets[0] = scan.offsets[1]
    assert float(scan.min_offset_separation().detach()) == pytest.approx(0.0)
    collapsed = torch.tensor([0.0, 0.0, 1.0], dtype=torch.float64)
    assert float(min_offset_separation(collapsed)) == pytest.approx(0.0)

python -m pytest packages/omnibias-torch/tests/test_scan.py -k test_min_offset_separation_on_learnable_offsets -q

Fields & methods

offsetsfield

offsets: Tensor

Declared Tensor field on BiasScan. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

Regression example · dtype none uses framework default

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_none_uses_framework_default() -> None:
    prev = torch.get_default_dtype()
    try:
        torch.set_default_dtype(torch.float64)
        scan = BiasScan(1, BankSpec.uniform(-1.0, 1.0, 3), dtype=None)
        assert scan.offsets.dtype == torch.float64
    finally:
        torch.set_default_dtype(prev)

python -m pytest packages/omnibias-torch/tests/test_scan.py -k test_dtype_none_uses_framework_default -q

View source ↗

scalesfield

scales: Tensor

Declared Tensor field on BiasScan. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

pool_tapsfield

pool_taps: Tensor

Declared Tensor field on BiasScan. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

__init__method

__init__(self, num_channels: int, bank: BankSpec, *, template: MultiPackSpec | OpName | str | None=None, op: OpName | str | None=None, base: str | ActivationSpec[Tensor]='tanh', learnable_offsets: bool=True, learnable_scales: bool=False, readout: Readout='response', gamma: float=8.0, derivative_order: int=1, gap: float=1.0, dtype: torch.dtype | None=None) -> None

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

Parameters

Name	Type & default	Contract
num_channelspositional-or-keyword	intRequired	See the operation contract and implementation below.
bankpositional-or-keyword	BankSpecRequired	See the operation contract and implementation below.
templatekeyword-only	MultiPackSpec | OpName | str | NoneDefault: None	See the operation contract and implementation below.
opkeyword-only	OpName | str | NoneDefault: None	See the operation contract and implementation below.
basekeyword-only	str | ActivationSpec[Tensor]Default: 'tanh'	See the operation contract and implementation below.
learnable_offsetskeyword-only	boolDefault: True	See the operation contract and implementation below.
learnable_scaleskeyword-only	boolDefault: False	See the operation contract and implementation below.
readoutkeyword-only	ReadoutDefault: 'response'	See the operation contract and implementation below.
gammakeyword-only	floatDefault: 8.0	See the operation contract and implementation below.
derivative_orderkeyword-only	intDefault: 1	See the operation contract and implementation below.
gapkeyword-only	floatDefault: 1.0	See the operation contract and implementation below.
dtypekeyword-only	torch.dtype | NoneDefault: None	See the operation contract and implementation below.

Returns: None

Validation & failure cases

• When num_channels < 1, raises ValueError(f'num_channels must be >= 1, got {num_channels}').
• When readout not in ('response', 'pooled', 'argmax'), raises ValueError(f'unknown readout {readout!r}').
• When op is not None, raises ValueError('pass only one of template= or op=; they name the same role').

Executed regression · dtype none uses framework default

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_dtype_none_uses_framework_default() -> None:
    prev = torch.get_default_dtype()
    try:
        torch.set_default_dtype(torch.float64)
        scan = BiasScan(1, BankSpec.uniform(-1.0, 1.0, 3), dtype=None)
        assert scan.offsets.dtype == torch.float64
    finally:
        torch.set_default_dtype(prev)

python -m pytest "packages/omnibias-torch/tests/test_scan.py::test_dtype_none_uses_framework_default" -q

Method implementation

    def __init__(
        self,
        num_channels: int,
        bank: BankSpec,
        *,
        template: MultiPackSpec | OpName | str | None = None,
        op: OpName | str | None = None,
        base: str | ActivationSpec[Tensor] = "tanh",
        learnable_offsets: bool = True,
        learnable_scales: bool = False,
        readout: Readout = "response",
        gamma: float = 8.0,
        derivative_order: int = 1,
        gap: float = 1.0,
        dtype: torch.dtype | None = None,
    ) -> None:
        super().__init__()
        if num_channels < 1:
            raise ValueError(f"num_channels must be >= 1, got {num_channels}")
        if readout not in ("response", "pooled", "argmax"):
            raise ValueError(f"unknown readout {readout!r}")
        self.num_channels = int(num_channels)
        self.readout: Readout = readout
        self.gamma = float(gamma)
        self.gap = float(gap)
        self.act_spec = base if isinstance(base, ActivationSpec) else get_activation(base)
        if isinstance(template, MultiPackSpec):
            if op is not None:
                raise ValueError("pass only one of template= or op=; they name the same role")
            self.template: MultiPackSpec | str = template
        else:
            role = resolve_scan_role(template, op)
            self.template = template_from_op(
                role, derivative_order=derivative_order, gap=gap
            )
        dt = torch.get_default_dtype() if dtype is None else dtype
        off0 = torch.tensor(list(bank.offsets), dtype=dt)
        sc0 = torch.tensor(list(bank.scales), dtype=dt)
        if learnable_offsets:
            self.offsets = nn.Parameter(off0)
        else:
            self.register_buffer("offsets", off0)
        if learnable_scales:
            self.scales = nn.Parameter(sc0)
        else:
            self.register_buffer("scales", sc0)
        taps0 = torch.full((off0.numel(),), 1.0 / max(off0.numel(), 1), dtype=dt)
        self.pool_taps = nn.Parameter(taps0)

View source ↗

min_offset_separationmethod

min_offset_separation(self) -> Tensor

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

Returns: Tensor

Executed regression · min offset separation on learnable offsets

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_min_offset_separation_on_learnable_offsets() -> None:
    torch.set_default_dtype(torch.float64)
    bank = BankSpec.uniform(-1.0, 1.0, 5)
    scan = _scan(bank=bank, learnable_offsets=True)
    assert float(scan.min_offset_separation().detach()) == pytest.approx(0.5)
    with torch.no_grad():
        scan.offsets[0] = scan.offsets[1]
    assert float(scan.min_offset_separation().detach()) == pytest.approx(0.0)
    collapsed = torch.tensor([0.0, 0.0, 1.0], dtype=torch.float64)
    assert float(min_offset_separation(collapsed)) == pytest.approx(0.0)

python -m pytest "packages/omnibias-torch/tests/test_scan.py::test_min_offset_separation_on_learnable_offsets" -q

Method implementation

    def min_offset_separation(self) -> Tensor:
        return min_offset_separation(self.offsets)

View source ↗

forwardmethod

forward(self, z: Tensor) -> Tensor

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

Parameters

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

Returns: Tensor

Validation & failure cases

• When z.shape[-1] != self.num_channels, raises ValueError(f'expected z[..., {self.num_channels}], got shape {tuple(z.shape)}').
• When self.readout == 'pooled', raises ValueError('pooled readout shape mismatch').

Executed regression · band and integral templates run

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_band_and_integral_templates_run() -> None:
    torch.set_default_dtype(torch.float64)
    bank = BankSpec.uniform(-1.0, 1.0, 3)
    z = torch.zeros(2, 1, dtype=torch.float64)
    band = BiasScan(1, bank, template="band", learnable_offsets=False, dtype=torch.float64)
    integ = BiasScan(
        1, bank, template="integral", base="tanh", learnable_offsets=False, dtype=torch.float64
    )
    assert band(z).shape == (2, 1, 3)
    assert integ(z).shape == (2, 1, 3)
    with pytest.raises(ValueError, match="unknown"):
        template_from_op("not-an-op")

python -m pytest "packages/omnibias-torch/tests/test_scan.py::test_band_and_integral_templates_run" -q

Method implementation

    def forward(self, z: Tensor) -> Tensor:
        if z.shape[-1] != self.num_channels:
            raise ValueError(
                f"expected z[..., {self.num_channels}], got shape {tuple(z.shape)}"
            )
        resp = scan_response(
            z, self.offsets, self.scales, self.template, self.act_spec, gap=self.gap
        )
        if self.readout == "response":
            return resp
        if self.readout == "pooled":
            taps = self.pool_taps.to(dtype=resp.dtype, device=resp.device)
            if resp.shape[-1] == taps.numel():
                return (resp * taps).sum(dim=-1)
            if resp.ndim >= 2 and resp.shape[-2] == taps.numel():
                pooled_m = (resp * taps.reshape(-1, 1)).sum(dim=-2)
                return pooled_m.mean(dim=-1)
            raise ValueError("pooled readout shape mismatch")
        loc = soft_argmax_offset(resp, self.offsets, gamma=self.gamma)
        if loc.ndim == z.ndim:
            return loc
        return loc.mean(dim=-1)

View source ↗

num_channelsattribute

num_channels = int(num_channels)

Trailing channel count of z.

View source ↗

readoutattribute

readout: Readout = readout

response -> (..., C, M); pooled / argmax -> (..., C).

View source ↗

gammaattribute

gamma = float(gamma)

Soft-argmax sharpness (not a collapse parameter).

View source ↗

gapattribute

gap = float(gap)

Public instance attribute assigned by __init__. The source assignment is gap = float(gap); it describes instance state, not a literal constructor default.

View source ↗

act_specattribute

act_spec = base if isinstance(base, ActivationSpec) else get_activation(base)

Public instance attribute assigned by __init__. The source assignment is act_spec = base if isinstance(base, ActivationSpec) else get_activation(base); it describes instance state, not a literal constructor default.

View source ↗

templateattribute

template: MultiPackSpec | str = template

~omnibias.core.multipack.MultiPackSpec or an operator role name.

View source ↗

Read the implementation

class BiasScan(nn.Module):
    """Channel-wise bias scan.

    Parameters
    ----------
    num_channels:
        Trailing channel count of ``z``.
    bank:
        Offset / scale bank.
    template:
        :class:`~omnibias.core.multipack.MultiPackSpec` or an operator role name.
    op:
        01-13 catalog alias for ``template``. Pass only one of ``template``
        or ``op``; they name the same six-role family, not a seventh role.
    readout:
        ``response`` -> ``(..., C, M)``; ``pooled`` / ``argmax`` -> ``(..., C)``.
    gamma:
        Soft-argmax sharpness (not a collapse parameter).
    """

    offsets: Tensor
    scales: Tensor
    pool_taps: Tensor

    def __init__(
        self,
        num_channels: int,
        bank: BankSpec,
        *,
        template: MultiPackSpec | OpName | str | None = None,
        op: OpName | str | None = None,
        base: str | ActivationSpec[Tensor] = "tanh",
        learnable_offsets: bool = True,
        learnable_scales: bool = False,
        readout: Readout = "response",
        gamma: float = 8.0,
        derivative_order: int = 1,
        gap: float = 1.0,
        dtype: torch.dtype | None = None,
    ) -> None:
        super().__init__()
        if num_channels < 1:
            raise ValueError(f"num_channels must be >= 1, got {num_channels}")
        if readout not in ("response", "pooled", "argmax"):
            raise ValueError(f"unknown readout {readout!r}")
        self.num_channels = int(num_channels)
        self.readout: Readout = readout
        self.gamma = float(gamma)
        self.gap = float(gap)
        self.act_spec = base if isinstance(base, ActivationSpec) else get_activation(base)
        if isinstance(template, MultiPackSpec):
            if op is not None:
                raise ValueError("pass only one of template= or op=; they name the same role")
            self.template: MultiPackSpec | str = template
        else:
            role = resolve_scan_role(template, op)
            self.template = template_from_op(
                role, derivative_order=derivative_order, gap=gap
            )
        dt = torch.get_default_dtype() if dtype is None else dtype
        off0 = torch.tensor(list(bank.offsets), dtype=dt)
        sc0 = torch.tensor(list(bank.scales), dtype=dt)
        if learnable_offsets:
            self.offsets = nn.Parameter(off0)
        else:
            self.register_buffer("offsets", off0)
        if learnable_scales:
            self.scales = nn.Parameter(sc0)
        else:
            self.register_buffer("scales", sc0)
        taps0 = torch.full((off0.numel(),), 1.0 / max(off0.numel(), 1), dtype=dt)
        self.pool_taps = nn.Parameter(taps0)

    def min_offset_separation(self) -> Tensor:
        return min_offset_separation(self.offsets)

    def forward(self, z: Tensor) -> Tensor:
        if z.shape[-1] != self.num_channels:
            raise ValueError(
                f"expected z[..., {self.num_channels}], got shape {tuple(z.shape)}"
            )
        resp = scan_response(
            z, self.offsets, self.scales, self.template, self.act_spec, gap=self.gap
        )
        if self.readout == "response":
            return resp
        if self.readout == "pooled":
            taps = self.pool_taps.to(dtype=resp.dtype, device=resp.device)
            if resp.shape[-1] == taps.numel():
                return (resp * taps).sum(dim=-1)
            if resp.ndim >= 2 and resp.shape[-2] == taps.numel():
                pooled_m = (resp * taps.reshape(-1, 1)).sum(dim=-2)
                return pooled_m.mean(dim=-1)
            raise ValueError("pooled readout shape mismatch")
        loc = soft_argmax_offset(resp, self.offsets, gamma=self.gamma)
        if loc.ndim == z.ndim:
            return loc
        return loc.mean(dim=-1)

Related APIsomnibias.core.scan.resolve_scan_role↗omnibias.torch.get_activation↗omnibias.torch.scan.min_offset_separation↗omnibias.torch.scan_response↗omnibias.torch.soft_argmax_offset↗omnibias.torch.scan.template_from_op↗

Other public import paths

• omnibias.torch.BiasScan
• omnibias.torch.scan.BiasScan

FUNCTION

min_offset_separation

omnibias.torch.scan.min_offset_separationSource ↗

min_offset_separation(offsets: Tensor) -> Tensor

Minimum adjacent spacing of offsets after sorting.

Parameters

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

ReturnsTensor

Executed regression · min offset separation on learnable offsets

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_min_offset_separation_on_learnable_offsets() -> None:
    torch.set_default_dtype(torch.float64)
    bank = BankSpec.uniform(-1.0, 1.0, 5)
    scan = _scan(bank=bank, learnable_offsets=True)
    assert float(scan.min_offset_separation().detach()) == pytest.approx(0.5)
    with torch.no_grad():
        scan.offsets[0] = scan.offsets[1]
    assert float(scan.min_offset_separation().detach()) == pytest.approx(0.0)
    collapsed = torch.tensor([0.0, 0.0, 1.0], dtype=torch.float64)
    assert float(min_offset_separation(collapsed)) == pytest.approx(0.0)

python -m pytest "packages/omnibias-torch/tests/test_scan.py::test_min_offset_separation_on_learnable_offsets" -q

Read the implementation

def min_offset_separation(offsets: Tensor) -> Tensor:
    """Minimum adjacent spacing of ``offsets`` after sorting."""
    ordered, _ = torch.sort(offsets.reshape(-1))
    if ordered.numel() < 2:
        return offsets.new_tensor(float("inf"))
    return (ordered[1:] - ordered[:-1]).min()

FUNCTION

scan_response

omnibias.torch.scan_responseSource ↗

scan_response(z: Tensor, offsets: Tensor, scales: Tensor, spec: MultiPackSpec | str, base: ActivationSpec[Tensor], *, gap: float=1.0) -> Tensor

Evaluate the template at scale * (z + offset).

z is (..., C), offsets (M,), scales (S,). Returns (..., C, M) when S == 1, else (..., C, M, S).

Parameters

Name	Type & default	Contract
zpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
offsetspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
scalespositional-or-keyword	TensorRequired	See the operation contract and implementation below.
specpositional-or-keyword	MultiPackSpec | strRequired	See the operation contract and implementation below.
basepositional-or-keyword	ActivationSpec[Tensor]Required	See the operation contract and implementation below.
gapkeyword-only	floatDefault: 1.0	See the operation contract and implementation below.

ReturnsTensor

Explicit validation & failure cases

• When base.forward is None, raises NotImplementedError('band template needs ActivationSpec.forward').
• When base.integral is None, raises NotImplementedError(f'Activation {base.name!r} has no closed-form integral kernel').

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

Executed regression · band and integral templates run

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_band_and_integral_templates_run() -> None:
    torch.set_default_dtype(torch.float64)
    bank = BankSpec.uniform(-1.0, 1.0, 3)
    z = torch.zeros(2, 1, dtype=torch.float64)
    band = BiasScan(1, bank, template="band", learnable_offsets=False, dtype=torch.float64)
    integ = BiasScan(
        1, bank, template="integral", base="tanh", learnable_offsets=False, dtype=torch.float64
    )
    assert band(z).shape == (2, 1, 3)
    assert integ(z).shape == (2, 1, 3)
    with pytest.raises(ValueError, match="unknown"):
        template_from_op("not-an-op")

python -m pytest "packages/omnibias-torch/tests/test_scan.py::test_band_and_integral_templates_run" -q

Regression example · two interface soft argmax is biased visible

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_two_interface_soft_argmax_is_biased_visible() -> None:
    """Two peaks one spacing apart: soft-argmax collapses to 0; do not hide it."""
    torch.set_default_dtype(torch.float64)
    bank = BankSpec((-1.0, -0.5, 0.0, 0.5, 1.0))
    offsets = torch.tensor(bank.offsets, dtype=torch.float64)
    scales = torch.tensor(bank.scales, dtype=torch.float64)
    from omnibias.torch.activations.registry import get_activation

    spec = template_from_op("grad")
    base = get_activation("tanh")
    r1 = scan_response(
        torch.tensor([[0.4]], dtype=torch.float64), offsets, scales, spec, base
    )
    r2 = scan_response(
        torch.tensor([[-0.4]], dtype=torch.float64), offsets, scales, spec, base
    )
    combined = r1 + r2
    tau_star = float(soft_argmax_offset(combined.squeeze(), offsets, gamma=8.0))
    assert abs(tau_star) < 0.15
    assert abs(tau_star - 0.4) > 0.2
    assert abs(tau_star + 0.4) > 0.2

python -m pytest packages/omnibias-torch/tests/test_scan.py -k test_two_interface_soft_argmax_is_biased_visible -q

Read the implementation

def scan_response(
    z: Tensor,
    offsets: Tensor,
    scales: Tensor,
    spec: MultiPackSpec | str,
    base: ActivationSpec[Tensor],
    *,
    gap: float = 1.0,
) -> Tensor:
    """Evaluate the template at ``scale * (z + offset)``.

    ``z`` is ``(..., C)``, ``offsets`` ``(M,)``, ``scales`` ``(S,)``.
    Returns ``(..., C, M)`` when ``S == 1``, else ``(..., C, M, S)``.
    """
    off = offsets.reshape(-1)
    sc = scales.reshape(-1)
    bank = z.unsqueeze(-1) + off
    if sc.numel() == 1:
        u = sc.reshape(()) * bank
    else:
        u = bank.unsqueeze(-1) * sc
    if isinstance(spec, str):
        half = 0.5 * float(gap)
        if spec == "band":
            if base.forward is None:
                raise NotImplementedError("band template needs ActivationSpec.forward")
            out = base.forward(u + half) - base.forward(u - half)
        elif spec == "integral":
            if base.integral is None:
                raise NotImplementedError(
                    f"Activation {base.name!r} has no closed-form integral kernel"
                )
            out = base.integral(u + half) - base.integral(u - half)
        else:
            raise ValueError(f"unknown window template {spec!r}")
    else:
        orders = tuple(int(p.order) for p in spec.packs)
        means = torch.tensor([p.mean for p in spec.packs], dtype=u.dtype, device=u.device)
        weights = torch.tensor([p.weight for p in spec.packs], dtype=u.dtype, device=u.device)
        out = multipack_response(u, means, weights, orders, base)
    return out

Related APIsomnibias.torch.multipack_response↗

Other public import paths

• omnibias.torch.scan_response
• omnibias.torch.scan.scan_response

FUNCTION

soft_argmax_offset

omnibias.torch.soft_argmax_offsetSource ↗

soft_argmax_offset(response: Tensor, offsets: Tensor, *, gamma: float=8.0) -> Tensor

tau* = sum_j tau_j softmax(gamma R)_j over the offset axis.

response is (..., M) or (..., M, S) (offset axis is -1 or -2). Softmax is over offsets. gamma -> inf would be a hard argmax (temperature collapse), not bias collapse.

Parameters

Name	Type & default	Contract
responsepositional-or-keyword	TensorRequired	See the operation contract and implementation below.
offsetspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
gammakeyword-only	floatDefault: 8.0	See the operation contract and implementation below.

ReturnsTensor

Executed regression · soft argmax matches finite difference dtau dz

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_soft_argmax_matches_finite_difference_dtau_dz() -> None:
    torch.set_default_dtype(torch.float64)
    bank = BankSpec.uniform(-1.0, 1.0, 5)
    scan = _scan(bank=bank, readout="argmax")
    z = torch.tensor([[0.3]], dtype=torch.float64, requires_grad=True)
    tau = scan(z)
    tau.backward()
    assert z.grad is not None
    analytic = float(z.grad.item())
    eps = 1e-6
    with torch.no_grad():
        tp = float(scan(torch.tensor([[0.3 + eps]], dtype=torch.float64)))
        tm = float(scan(torch.tensor([[0.3 - eps]], dtype=torch.float64)))
    fd = (tp - tm) / (2.0 * eps)
    assert analytic == pytest.approx(fd, rel=1e-4, abs=1e-6)

python -m pytest "packages/omnibias-torch/tests/test_scan.py::test_soft_argmax_matches_finite_difference_dtau_dz" -q

Regression example · worked example matches spec

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_worked_example_matches_spec() -> None:
    """Spec §5: tanh grad bank at z=0.3, gamma=8, tau* ~ -0.283."""
    torch.set_default_dtype(torch.float64)
    bank = BankSpec((-1.0, -0.5, 0.0, 0.5, 1.0))
    z = torch.tensor([[0.3]], dtype=torch.float64)
    resp = _scan(bank=bank, readout="response")(z).squeeze()
    expected_u = z.item() + torch.tensor(bank.offsets, dtype=torch.float64)
    expected_r = 1.0 - torch.tanh(expected_u).pow(2)
    assert torch.equal(resp, expected_r)
    offs = torch.tensor(bank.offsets, dtype=torch.float64)
    tau_star = float(soft_argmax_offset(resp, offs, gamma=8.0))
    weights = torch.softmax(8.0 * resp, dim=0)
    expected_tau = float((weights * offs).sum())
    assert tau_star == expected_tau
    # Spec §5 quoted tau* ~ -0.283 from a rounded R-table; the exact
    # softmax of 1-tanh^2 is within a tenth of a spacing of -z = -0.3.
    assert abs(tau_star + 0.3) <= 0.1 * 0.5

python -m pytest packages/omnibias-torch/tests/test_scan.py -k test_worked_example_matches_spec -q

Read the implementation

def soft_argmax_offset(response: Tensor, offsets: Tensor, *, gamma: float = 8.0) -> Tensor:
    """``tau* = sum_j tau_j softmax(gamma R)_j`` over the offset axis.

    ``response`` is ``(..., M)`` or ``(..., M, S)`` (offset axis is -1 or -2).
    Softmax is over offsets. ``gamma -> inf`` would be a hard argmax
    (temperature collapse), not bias collapse.
    """
    off = offsets.reshape(-1).to(dtype=response.dtype, device=response.device)
    if response.shape[-1] == off.numel():
        logits = float(gamma) * response
        w = F.softmax(logits, dim=-1)
        return (w * off).sum(dim=-1)
    if response.ndim >= 2 and response.shape[-2] == off.numel():
        logits = float(gamma) * response
        w = F.softmax(logits, dim=-2)
        return (w * off.reshape(-1, 1)).sum(dim=-2)
    raise ValueError("response trailing shape does not match offsets")

Other public import paths

• omnibias.torch.soft_argmax_offset
• omnibias.torch.scan.soft_argmax_offset

FUNCTION

template_from_op

omnibias.torch.scan.template_from_opSource ↗

template_from_op(op: OpName | str, *, derivative_order: int=1, gap: float=1.0) -> MultiPackSpec | str

Map an operator role to a MultiPackSpec, or keep window ops as tags.

Parameters

Name	Type & default	Contract
oppositional-or-keyword	OpName | strRequired	See the operation contract and implementation below.
derivative_orderkeyword-only	intDefault: 1	See the operation contract and implementation below.
gapkeyword-only	floatDefault: 1.0	See the operation contract and implementation below.

ReturnsMultiPackSpec | str

Explicit validation & failure cases

• When int(derivative_order) < 0, raises ValueError('derivative_order must be >= 0').
• When gap <= 0.0, raises ValueError('gap must be positive').

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

Executed regression · dtype none uses framework default

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_dtype_none_uses_framework_default() -> None:
    prev = torch.get_default_dtype()
    try:
        torch.set_default_dtype(torch.float64)
        scan = BiasScan(1, BankSpec.uniform(-1.0, 1.0, 3), dtype=None)
        assert scan.offsets.dtype == torch.float64
    finally:
        torch.set_default_dtype(prev)

python -m pytest "packages/omnibias-torch/tests/test_scan.py::test_dtype_none_uses_framework_default" -q

Regression example · unknown template raises

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_unknown_template_raises() -> None:
    with pytest.raises(ValueError, match="unknown"):
        from omnibias.torch.scan import template_from_op

        template_from_op("nope")

python -m pytest tests/test_scan_parity.py -k test_unknown_template_raises -q

Read the implementation

def template_from_op(
    op: OpName | str,
    *,
    derivative_order: int = 1,
    gap: float = 1.0,
) -> MultiPackSpec | str:
    """Map an operator role to a :class:`MultiPackSpec`, or keep window ops as tags."""
    name = str(op).lower()
    if name in _COLLAPSE_ORDER:
        return MultiPackSpec((PackSpec(_COLLAPSE_ORDER[name], 0.0),))
    if name == "derivative":
        if int(derivative_order) < 0:
            raise ValueError("derivative_order must be >= 0")
        return MultiPackSpec((PackSpec(int(derivative_order), 0.0),))
    if name in ("band", "integral"):
        if gap <= 0.0:
            raise ValueError("gap must be positive")
        return name
    raise ValueError(f"unknown template op {op!r}")

Related APIsomnibias.core.MultiPackSpec↗omnibias.core.PackSpec↗

← Package overviewRunnable learning examples →