omnibias.torch.unit
PyTorch activation derivatives, trainable operator layers, network jets, and direct differential-operator contractions.
import omnibias.torch.unitModule contract & mathematical conventions
OperatorMultiBiasUnit (OMBU): the trainable scalar-operator primitive.
A K-bias multi-bias unit applied per channel:
f_K(z; b, s) = sum_{k=1}^{K} s_k * sigma(z + b_k)
with K learnable biases b_1..b_K and K signs s_1..s_K per channel,
on top of a fixed base activation sigma.
Two operating regimes:
1. **Identity-nested (Lemma 1).** With tied biases and signs summing to one, the unit reduces by the stated algebraic identity to sigma(z + b). This is the default initialisation for any K, so a freshly-instantiated OMBU is a drop-in replacement for sigma at step zero.
2. **Bias-collapse derivative tower (Lemma collapse).** With biases on the forward-difference (or central-difference) stencil and rescaled signs, the unit converges to sigma^(K-1)(z + b_mean) as the bias spread tends to zero. The closed-form derivative is available via analytic_derivative, which uses the per-activation fast-path kernel (Eulerian polynomial for sigmoid, Hermite for Gaussian, etc.).
OMBU keeps biases and signs as free parameters; the operator-typed
omnibias.blocks.OperatorBlock is the wrapper that picks K and
calls either the literal forward or analytic_derivative depending
on the operator role.
Terminology: this delta -> 0 limit is *the* **founding bias collapse**
(many biases collapse onto one value, yielding a derivative). It is distinct
from **temperature collapse**, the beta -> inf feasibility penalty in
omnibias.convex (and omnibias.control / omnibias.routing), which
sharpens one constraint into a 0/1 feasibility step -- an indicator, not a derivative. See
docs/derivatives.md.
Numerical scope: shared coefficients and formulas do not establish universal floating-point identity across backends, devices, compiler transforms, or reduction orders. Compare the chosen path with dtype-appropriate tolerances; even float64 is not a general bitwise-equality guarantee. Algebraic initialization identities likewise need a numerical check for the selected coefficients and reduction order.
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.
OperatorMultiBiasUnit
omnibias.torch.OMBUSource ↗class OperatorMultiBiasUnit(nn.Module)The trainable scalar-operator primitive.
Parameters
num_channels : int Number of independent channels. Each channel has its own K biases (and optionally its own K signs).
K : int, default 2 Number of bias terms per channel. K=1 recovers the standard single-bias activation.
base : str or ActivationSpec, default "sigmoid" Base activation. May be a string from omnibias.activations.list_activations or a custom spec.
init_bias : float, default 0.0 Initial value for every bias entry (Lemma 1: tied biases recover sigma(z + init_bias) bit-identically).
init_delta : float, default 0.0 If non-zero, initialises biases on the central-difference stencil b + (k - (K+1)/2) * init_delta. Use 0.0 for strict Lemma 1 identity-nesting; use a small positive value to break the symmetry at init for non-monotone shapes.
learnable_biases : bool, default True If False, biases are registered as a buffer (frozen).
learnable_signs : bool, default False If True, signs are a learnable parameter (mb-K-lrn-s in the paper). If False, signs are a frozen buffer initialised by identity_init_signs.
Parameters
| Name | Type & default | Contract |
|---|---|---|
num_channelspositional-or-keyword | intRequired | int Number of independent channels. Each channel has its own K biases (and optionally its own K signs). |
Kpositional-or-keyword | intDefault: 2 | int, default 2 Number of bias terms per channel. K=1 recovers the standard single-bias activation. |
basepositional-or-keyword | str | ActivationSpec[Tensor]Default: 'sigmoid' | str or ActivationSpec, default "sigmoid" Base activation. May be a string from omnibias.activations.list_activations or a custom spec. |
init_biaspositional-or-keyword | floatDefault: 0.0 | float, default 0.0 Initial value for every bias entry (Lemma 1: tied biases recover sigma(z + init_bias) bit-identically). |
init_deltapositional-or-keyword | floatDefault: 0.0 | float, default 0.0 If non-zero, initialises biases on the central-difference stencil b + (k - (K+1)/2) * init_delta. Use 0.0 for strict Lemma 1 identity-nesting; use a small positive value to break the symmetry at init for non-monotone shapes. |
learnable_biasespositional-or-keyword | boolDefault: True | bool, default True If False, biases are registered as a buffer (frozen). |
learnable_signspositional-or-keyword | boolDefault: False | bool, default False If True, signs are a learnable parameter (mb-K-lrn-s in the paper). If False, signs are a frozen buffer initialised by identity_init_signs. |
biases_initpositional-or-keyword | Tensor | NoneDefault: None | See the operation contract and implementation below. |
signs_initpositional-or-keyword | Tensor | NoneDefault: None | See the operation contract and implementation below. |
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.
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" -qRegression example · invalid input shape 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.
def test_invalid_input_shape_raises() -> None:
ombu = OperatorMultiBiasUnit(num_channels=4, K=2, base="sigmoid")
with pytest.raises(ValueError):
ombu(torch.randn(3, 5))python -m pytest packages/omnibias-torch/tests/test_unit.py -k test_invalid_input_shape_raises -qFields & methods
specfield
spec: ActivationSpec[Tensor]Declared ActivationSpec[Tensor] field on OperatorMultiBiasUnit. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.
Regression example · construction defaults
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.
def test_construction_defaults() -> None:
ombu = OperatorMultiBiasUnit(num_channels=4, K=2, base="sigmoid")
assert ombu.num_channels == 4
assert ombu.K == 2
assert ombu.spec.name == "sigmoid"
assert ombu.is_identity_nested
assert ombu.biases.shape == (4, 2)
assert ombu.signs.shape == (4, 2)python -m pytest packages/omnibias-torch/tests/test_unit.py -k test_construction_defaults -qbiasesfield
biases: TensorDeclared Tensor field on OperatorMultiBiasUnit. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.
Regression example · ombu uses 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.
def test_ombu_uses_default_dtype(torch_default_float64) -> None:
ombu = OMBU(num_channels=4, K=2, base="tanh")
assert ombu.biases.dtype == torch.float64
assert ombu.signs.dtype == torch.float64
out = ombu(torch.zeros(8, 4))
assert out.dtype == torch.float64python -m pytest packages/omnibias-torch/tests/test_audit_regressions.py -k test_ombu_uses_default_dtype -qsignsfield
signs: TensorDeclared Tensor field on OperatorMultiBiasUnit. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.
Regression example · buffer signs stay buffer
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.
def test_buffer_signs_stay_buffer() -> None:
ombu = OperatorMultiBiasUnit(num_channels=2, K=2, base="sigmoid", learnable_signs=False)
assert not isinstance(ombu.signs, torch.nn.Parameter)python -m pytest packages/omnibias-torch/tests/test_unit.py -k test_buffer_signs_stay_buffer -q__init__method
__init__(self, num_channels: int, K: int=2, base: str | ActivationSpec[Tensor]='sigmoid', init_bias: float=0.0, init_delta: float=0.0, learnable_biases: bool=True, learnable_signs: bool=False, biases_init: Tensor | None=None, signs_init: Tensor | None=None) -> NoneSee 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. |
Kpositional-or-keyword | intDefault: 2 | See the operation contract and implementation below. |
basepositional-or-keyword | str | ActivationSpec[Tensor]Default: 'sigmoid' | See the operation contract and implementation below. |
init_biaspositional-or-keyword | floatDefault: 0.0 | See the operation contract and implementation below. |
init_deltapositional-or-keyword | floatDefault: 0.0 | See the operation contract and implementation below. |
learnable_biasespositional-or-keyword | boolDefault: True | See the operation contract and implementation below. |
learnable_signspositional-or-keyword | boolDefault: False | See the operation contract and implementation below. |
biases_initpositional-or-keyword | Tensor | NoneDefault: None | See the operation contract and implementation below. |
signs_initpositional-or-keyword | Tensor | NoneDefault: None | See the operation contract and implementation below. |
Returns: None
Validation & failure cases
- When
num_channels < 1, raisesValueError(f'num_channels must be >= 1, got {num_channels}'). - When
K < 1, raisesValueError(f'K must be >= 1, got {K}'). - When
init_delta < 0, raisesValueError(f'init_delta must be >= 0, got {init_delta}'). - When
biases_init.shape != (num_channels, K), raisesValueError(f'biases_init shape {tuple(biases_init.shape)} does not match ({num_channels}, {K}).'). - When
signs_init.shape != (num_channels, K), raisesValueError(f'signs_init shape {tuple(signs_init.shape)} does not match ({num_channels}, {K}) (or ({K},) for broadcast).').
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.
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" -qMethod implementation
def __init__(
self,
num_channels: int,
K: int = 2,
base: str | ActivationSpec[Tensor] = "sigmoid",
init_bias: float = 0.0,
init_delta: float = 0.0,
learnable_biases: bool = True,
learnable_signs: bool = False,
biases_init: Tensor | None = None,
signs_init: Tensor | None = None,
) -> None:
super().__init__()
if num_channels < 1:
raise ValueError(f"num_channels must be >= 1, got {num_channels}")
if K < 1:
raise ValueError(f"K must be >= 1, got {K}")
if init_delta < 0:
raise ValueError(f"init_delta must be >= 0, got {init_delta}")
self.num_channels = num_channels
self.K = K
self.spec = base if isinstance(base, ActivationSpec) else get_activation(base)
if biases_init is None:
biases_init = identity_init_biases(num_channels, K, bias_value=init_bias)
if init_delta > 0 and K > 1:
offsets = central_bias_offsets(K, init_delta).unsqueeze(0) # (1, K)
biases_init = biases_init + offsets
else:
biases_init = torch.as_tensor(biases_init, dtype=torch.get_default_dtype())
if biases_init.shape != (num_channels, K):
raise ValueError(
f"biases_init shape {tuple(biases_init.shape)} does not match "
f"({num_channels}, {K})."
)
if signs_init is None:
signs_init = identity_init_signs(num_channels, K)
else:
signs_init = torch.as_tensor(signs_init, dtype=torch.get_default_dtype())
if signs_init.shape == (K,):
signs_init = signs_init.unsqueeze(0).expand(num_channels, K).contiguous()
elif signs_init.shape != (num_channels, K):
raise ValueError(
f"signs_init shape {tuple(signs_init.shape)} does not match "
f"({num_channels}, {K}) (or ({K},) for broadcast)."
)
if learnable_biases:
self.biases = nn.Parameter(biases_init)
else:
self.register_buffer("biases", biases_init)
if learnable_signs:
self.signs = nn.Parameter(signs_init)
else:
self.register_buffer("signs", signs_init)forwardmethod
forward(self, z: Tensor) -> TensorLiteral multi-bias forward.
Computes sum_k s_k * sigma(z + b_k) channel-wise.
z must have shape (..., num_channels); the return value
has the same shape.
Parameters
| Name | Type & default | Contract |
|---|---|---|
zpositional-or-keyword | TensorRequired | See the operation contract and implementation below. |
Returns: Tensor
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.
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" -qMethod implementation
def forward(self, z: Tensor) -> Tensor:
"""Literal multi-bias forward.
Computes ``sum_k s_k * sigma(z + b_k)`` channel-wise.
``z`` must have shape ``(..., num_channels)``; the return value
has the same shape.
"""
return multibias_literal_forward(z, self.biases, self.signs, self.spec.forward)analytic_derivativemethod
analytic_derivative(self, z: Tensor, order: int | None=None) -> TensorClosed-form sigma^(order)(z + bias_mean) for activations
whose derivative tower has a known polynomial form (sigmoid,
tanh, gaussian, ...).
order defaults to K - 1 (the standard bias-collapse limit
order for this unit).
Raises
NotImplementedError If the spec has no fast-path kernel.
Parameters
| Name | Type & default | Contract |
|---|---|---|
zpositional-or-keyword | TensorRequired | See the operation contract and implementation below. |
orderpositional-or-keyword | int | NoneDefault: None | See the operation contract and implementation below. |
Returns: Tensor
Validation & failure cases
- When
order < 0, raisesValueError(f'order must be >= 0, got {order}'). - When
self.spec.fastpath is None, raisesNotImplementedError(f'Activation {self.spec.name!r} has no closed-form derivative kernel; use forward() instead, or pick a base from the smooth/proximal families.').
Executed regression · cmbconv1d 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.
def test_cmbconv1d_shape() -> None:
layer = cmbConv1d(2, 4, kernel_size=3, padding=1, op="grad", base="tanh")
out = layer(torch.randn(1, 2, 16))
assert out.shape == (1, 4, 16)python -m pytest "packages/omnibias-torch/tests/test_blocks.py::test_cmbconv1d_shape" -qMethod implementation
def analytic_derivative(self, z: Tensor, order: int | None = None) -> Tensor:
"""Closed-form ``sigma^(order)(z + bias_mean)`` for activations
whose derivative tower has a known polynomial form (sigmoid,
tanh, gaussian, ...).
``order`` defaults to ``K - 1`` (the standard bias-collapse limit
order for this unit).
Raises
------
NotImplementedError
If the spec has no fast-path kernel.
"""
if order is None:
order = self.K - 1
if order < 0:
raise ValueError(f"order must be >= 0, got {order}")
if self.spec.fastpath is None:
raise NotImplementedError(
f"Activation {self.spec.name!r} has no closed-form derivative kernel; "
"use forward() instead, or pick a base from the smooth/proximal families."
)
b_mean = self.biases.mean(dim=-1) # (num_channels,)
out: Tensor = self.spec.fastpath(z + b_mean, order)
return outanalytic_integralmethod
analytic_integral(self, z: Tensor, *, strict_window: bool=False, normalize: bool=False, small_width_threshold: float=0.0001) -> TensorClosed-form sign-weighted antiderivative sum.
If S' = sigma, the default path evaluates
sum_k s_k S(z + b_k). With strict_window=True and K=2,
the two stored bias parameters are interpreted as (center,
raw_width) and converted to an ordered low-to-high bias window.
Parameters
| Name | Type & default | Contract |
|---|---|---|
zpositional-or-keyword | TensorRequired | See the operation contract and implementation below. |
strict_windowkeyword-only | boolDefault: False | 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. |
Returns: Tensor
Validation & failure cases
- When
self.K != 2, raisesValueError('strict integral windows require K=2.').
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.
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" -qMethod implementation
def analytic_integral(
self,
z: Tensor,
*,
strict_window: bool = False,
normalize: bool = False,
small_width_threshold: float = 1e-4,
) -> Tensor:
"""Closed-form sign-weighted antiderivative sum.
If ``S' = sigma``, the default path evaluates
``sum_k s_k S(z + b_k)``. With ``strict_window=True`` and ``K=2``,
the two stored bias parameters are interpreted as ``(center,
raw_width)`` and converted to an ordered low-to-high bias window.
"""
if strict_window:
if self.K != 2:
raise ValueError("strict integral windows require K=2.")
center = self.biases[:, 0]
width = F.softplus(self.biases[:, 1])
use_taylor = not isinstance(self.signs, nn.Parameter)
return multibias_integral_window_forward(
z,
center,
width,
self.signs,
self.spec,
normalize=normalize,
small_width_threshold=small_width_threshold,
use_small_width_taylor=use_taylor,
)
return multibias_integral_forward(z, self.biases, self.signs, self.spec)bias_spreadproperty
bias_spread(self) -> TensorPer-channel max(b_k) - min(b_k); near zero in the collapse regime.
Returns: Tensor
Executed regression · construction defaults
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.
def test_construction_defaults() -> None:
ombu = OperatorMultiBiasUnit(num_channels=4, K=2, base="sigmoid")
assert ombu.num_channels == 4
assert ombu.K == 2
assert ombu.spec.name == "sigmoid"
assert ombu.is_identity_nested
assert ombu.biases.shape == (4, 2)
assert ombu.signs.shape == (4, 2)python -m pytest "packages/omnibias-torch/tests/test_unit.py::test_construction_defaults" -qMethod implementation
def bias_spread(self) -> Tensor:
"""Per-channel ``max(b_k) - min(b_k)``; near zero in the collapse regime."""
with torch.no_grad():
return self.biases.max(dim=-1).values - self.biases.min(dim=-1).valuesis_identity_nestedproperty
is_identity_nested(self) -> boolTrue iff biases are tied and signs sum to one (Lemma 1).
Returns: bool
Executed regression · construction defaults
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.
def test_construction_defaults() -> None:
ombu = OperatorMultiBiasUnit(num_channels=4, K=2, base="sigmoid")
assert ombu.num_channels == 4
assert ombu.K == 2
assert ombu.spec.name == "sigmoid"
assert ombu.is_identity_nested
assert ombu.biases.shape == (4, 2)
assert ombu.signs.shape == (4, 2)python -m pytest "packages/omnibias-torch/tests/test_unit.py::test_construction_defaults" -qMethod implementation
def is_identity_nested(self) -> bool:
"""True iff biases are tied and signs sum to one (Lemma 1)."""
with torch.no_grad():
tied = bool((self.bias_spread == 0).all().item())
sums_to_one = bool(((self.signs.sum(dim=-1) - 1.0).abs() < 1e-6).all().item())
return tied and sums_to_oneextra_reprmethod
extra_repr(self) -> strSee the class contract above and the exact method implementation below.
Returns: str
Method implementation
def extra_repr(self) -> str:
return (
f"num_channels={self.num_channels}, K={self.K}, base={self.spec.name!r}, "
f"learnable_biases={isinstance(self.biases, nn.Parameter)}, "
f"learnable_signs={isinstance(self.signs, nn.Parameter)}"
)num_channelsattribute
num_channels: int = num_channelsNumber of independent channels. Each channel has its own K biases (and optionally its own K signs).
Regression example · construction defaults
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.
def test_construction_defaults() -> None:
ombu = OperatorMultiBiasUnit(num_channels=4, K=2, base="sigmoid")
assert ombu.num_channels == 4
assert ombu.K == 2
assert ombu.spec.name == "sigmoid"
assert ombu.is_identity_nested
assert ombu.biases.shape == (4, 2)
assert ombu.signs.shape == (4, 2)python -m pytest packages/omnibias-torch/tests/test_unit.py -k test_construction_defaults -qKattribute
K: int = KNumber of bias terms per channel. K=1 recovers the standard single-bias activation.
Regression example · construction defaults
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.
def test_construction_defaults() -> None:
ombu = OperatorMultiBiasUnit(num_channels=4, K=2, base="sigmoid")
assert ombu.num_channels == 4
assert ombu.K == 2
assert ombu.spec.name == "sigmoid"
assert ombu.is_identity_nested
assert ombu.biases.shape == (4, 2)
assert ombu.signs.shape == (4, 2)python -m pytest packages/omnibias-torch/tests/test_unit.py -k test_construction_defaults -qRead the implementation
class OperatorMultiBiasUnit(nn.Module):
"""The trainable scalar-operator primitive.
Parameters
----------
num_channels : int
Number of independent channels. Each channel has its own K biases
(and optionally its own K signs).
K : int, default 2
Number of bias terms per channel. ``K=1`` recovers the standard
single-bias activation.
base : str or :class:`ActivationSpec`, default ``"sigmoid"``
Base activation. May be a string from
:func:`omnibias.activations.list_activations` or a custom spec.
init_bias : float, default 0.0
Initial value for every bias entry (Lemma 1: tied biases recover
``sigma(z + init_bias)`` bit-identically).
init_delta : float, default 0.0
If non-zero, initialises biases on the central-difference stencil
``b + (k - (K+1)/2) * init_delta``. Use ``0.0`` for strict Lemma 1
identity-nesting; use a small positive value to break the
symmetry at init for non-monotone shapes.
learnable_biases : bool, default True
If False, biases are registered as a buffer (frozen).
learnable_signs : bool, default False
If True, signs are a learnable parameter (mb-K-lrn-s in the paper).
If False, signs are a frozen buffer initialised by
:func:`identity_init_signs`.
"""
spec: ActivationSpec[Tensor]
biases: Tensor
signs: Tensor
def __init__(
self,
num_channels: int,
K: int = 2,
base: str | ActivationSpec[Tensor] = "sigmoid",
init_bias: float = 0.0,
init_delta: float = 0.0,
learnable_biases: bool = True,
learnable_signs: bool = False,
biases_init: Tensor | None = None,
signs_init: Tensor | None = None,
) -> None:
super().__init__()
if num_channels < 1:
raise ValueError(f"num_channels must be >= 1, got {num_channels}")
if K < 1:
raise ValueError(f"K must be >= 1, got {K}")
if init_delta < 0:
raise ValueError(f"init_delta must be >= 0, got {init_delta}")
self.num_channels = num_channels
self.K = K
self.spec = base if isinstance(base, ActivationSpec) else get_activation(base)
if biases_init is None:
biases_init = identity_init_biases(num_channels, K, bias_value=init_bias)
if init_delta > 0 and K > 1:
offsets = central_bias_offsets(K, init_delta).unsqueeze(0) # (1, K)
biases_init = biases_init + offsets
else:
biases_init = torch.as_tensor(biases_init, dtype=torch.get_default_dtype())
if biases_init.shape != (num_channels, K):
raise ValueError(
f"biases_init shape {tuple(biases_init.shape)} does not match "
f"({num_channels}, {K})."
)
if signs_init is None:
signs_init = identity_init_signs(num_channels, K)
else:
signs_init = torch.as_tensor(signs_init, dtype=torch.get_default_dtype())
if signs_init.shape == (K,):
signs_init = signs_init.unsqueeze(0).expand(num_channels, K).contiguous()
elif signs_init.shape != (num_channels, K):
raise ValueError(
f"signs_init shape {tuple(signs_init.shape)} does not match "
f"({num_channels}, {K}) (or ({K},) for broadcast)."
)
if learnable_biases:
self.biases = nn.Parameter(biases_init)
else:
self.register_buffer("biases", biases_init)
if learnable_signs:
self.signs = nn.Parameter(signs_init)
else:
self.register_buffer("signs", signs_init)
# ----- forward passes -----
def forward(self, z: Tensor) -> Tensor:
"""Literal multi-bias forward.
Computes ``sum_k s_k * sigma(z + b_k)`` channel-wise.
``z`` must have shape ``(..., num_channels)``; the return value
has the same shape.
"""
return multibias_literal_forward(z, self.biases, self.signs, self.spec.forward)
def analytic_derivative(self, z: Tensor, order: int | None = None) -> Tensor:
"""Closed-form ``sigma^(order)(z + bias_mean)`` for activations
whose derivative tower has a known polynomial form (sigmoid,
tanh, gaussian, ...).
``order`` defaults to ``K - 1`` (the standard bias-collapse limit
order for this unit).
Raises
------
NotImplementedError
If the spec has no fast-path kernel.
"""
if order is None:
order = self.K - 1
if order < 0:
raise ValueError(f"order must be >= 0, got {order}")
if self.spec.fastpath is None:
raise NotImplementedError(
f"Activation {self.spec.name!r} has no closed-form derivative kernel; "
"use forward() instead, or pick a base from the smooth/proximal families."
)
b_mean = self.biases.mean(dim=-1) # (num_channels,)
out: Tensor = self.spec.fastpath(z + b_mean, order)
return out
def analytic_integral(
self,
z: Tensor,
*,
strict_window: bool = False,
normalize: bool = False,
small_width_threshold: float = 1e-4,
) -> Tensor:
"""Closed-form sign-weighted antiderivative sum.
If ``S' = sigma``, the default path evaluates
``sum_k s_k S(z + b_k)``. With ``strict_window=True`` and ``K=2``,
the two stored bias parameters are interpreted as ``(center,
raw_width)`` and converted to an ordered low-to-high bias window.
"""
if strict_window:
if self.K != 2:
raise ValueError("strict integral windows require K=2.")
center = self.biases[:, 0]
width = F.softplus(self.biases[:, 1])
use_taylor = not isinstance(self.signs, nn.Parameter)
return multibias_integral_window_forward(
z,
center,
width,
self.signs,
self.spec,
normalize=normalize,
small_width_threshold=small_width_threshold,
use_small_width_taylor=use_taylor,
)
return multibias_integral_forward(z, self.biases, self.signs, self.spec)
# ----- introspection helpers -----
@property
def bias_spread(self) -> Tensor:
"""Per-channel ``max(b_k) - min(b_k)``; near zero in the collapse regime."""
with torch.no_grad():
return self.biases.max(dim=-1).values - self.biases.min(dim=-1).values
@property
def is_identity_nested(self) -> bool:
"""True iff biases are tied and signs sum to one (Lemma 1)."""
with torch.no_grad():
tied = bool((self.bias_spread == 0).all().item())
sums_to_one = bool(((self.signs.sum(dim=-1) - 1.0).abs() < 1e-6).all().item())
return tied and sums_to_one
def extra_repr(self) -> str:
return (
f"num_channels={self.num_channels}, K={self.K}, base={self.spec.name!r}, "
f"learnable_biases={isinstance(self.biases, nn.Parameter)}, "
f"learnable_signs={isinstance(self.signs, nn.Parameter)}"
)Other public import paths
omnibias.torch.OMBUomnibias.torch.OperatorMultiBiasUnitomnibias.torch.unit.OperatorMultiBiasUnit