omnibias.torch.multipack
PyTorch activation derivatives, trainable operator layers, network jets, and direct differential-operator contractions.
import omnibias.torch.multipackModule contract & mathematical conventions
Heterogeneous multi-pack Birkhoff unit (torch; theory 01-01).
Closed-form evaluation of
F(z) = sum_g c_g * sigma^(n_g)(z + mu_g)
via the activation fastpath. With share_means=True (default), packs that
share a mean in the ~omnibias.core.multipack.MultiPackSpec share one
mean Parameter. A registered tower provider evaluates the base once per
distinct mean; other activations retain the per-order fastpath fallback. Founding bias
collapse delta -> 0 only -- no temperature collapse.
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.
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 |
|---|---|
MultiPackSpec | omnibias.core.MultiPackSpec |
PackSpec | omnibias.core.PackSpec |
MultiPackUnit
omnibias.torch.BirkhoffOMBUSource ↗class MultiPackUnit(nn.Module)Channel-wise multi-pack Birkhoff unit.
Parameters
num_channels: Number of independent channels; z has shape (..., num_channels).
spec: ~omnibias.core.multipack.MultiPackSpec describing orders, initial means, and initial outer weights.
base: Activation name or ~omnibias.torch.activations.ActivationSpec.
learnable_means / learnable_weights: Whether means / outer weights are nn.Parameter or frozen buffers.
share_means: If true (default), packs with equal means in spec share one mean Parameter (one base evaluation per mean with a registered tower).
dtype: Parameter dtype; None resolves to torch.get_default_dtype().
Parameters
| Name | Type & default | Contract |
|---|---|---|
num_channelspositional-or-keyword | intRequired | Number of independent channels; z has shape (..., num_channels). |
specpositional-or-keyword | MultiPackSpecRequired | ~omnibias.core.multipack.MultiPackSpec describing orders, initial means, and initial outer weights. |
basekeyword-only | str | ActivationSpec[Tensor]Default: 'sigmoid' | Activation name or ~omnibias.torch.activations.ActivationSpec. |
learnable_meanskeyword-only | boolDefault: True | See the operation contract and implementation below. |
learnable_weightskeyword-only | boolDefault: True | See the operation contract and implementation below. |
share_meanskeyword-only | boolDefault: True | If true (default), packs with equal means in spec share one mean Parameter (one base evaluation per mean with a registered tower). |
dtypekeyword-only | torch.dtype | NoneDefault: None | Parameter dtype; None resolves to torch.get_default_dtype(). |
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.
def test_worked_example_matches_spec() -> None:
"""Spec §5 closed form at z=0 for tanh support."""
from omnibias.torch.multipack import MultiPackUnit
torch.set_default_dtype(torch.float64)
spec = MultiPackSpec((PackSpec(1, -0.5, 1.0), PackSpec(2, 0.5, 0.25)))
unit = MultiPackUnit(
1, spec, base="tanh", learnable_means=False, learnable_weights=False
)
z = torch.zeros(1, 1, dtype=torch.float64)
got = float(unit(z).item())
t1 = math.tanh(-0.5)
s1 = 1.0 - t1 * t1
t2 = math.tanh(0.5)
s2 = -2.0 * t2 * (1.0 - t2 * t2)
expected = 1.0 * s1 + 0.25 * s2
assert got == pytest.approx(
expected, rel=0, abs=4 * np.finfo(float).eps * max(abs(expected), 1.0)
)python -m pytest tests/test_multipack_parity.py -k test_worked_example_matches_spec -qRegression example · g2 collapse no closed form growth
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_g2_collapse_no_closed_form_growth() -> None:
"""Closed form stable as delta shrinks; FD error grows then breaks."""
from omnibias.torch.multipack import MultiPackUnit
torch.set_default_dtype(torch.float64)
spec = MultiPackSpec((PackSpec(1, -0.5, 1.0), PackSpec(2, 0.5, 0.25)))
unit = MultiPackUnit(
1, spec, base="tanh", learnable_means=False, learnable_weights=False
)
z0 = 0.0
closed = float(unit(torch.tensor([[z0]], dtype=torch.float64)).item())
deltas = [1e-2, 1e-3, 1e-4, 1e-5, 1e-6]
fd_errs = []
for delta in deltas:
acc = 0.0
for p in spec.packs:
offsets, signs = central_stencil_weights(p.order, delta)
for off, s in zip(offsets, signs, strict=True):
acc += p.weight * s * math.tanh(z0 + p.mean + off)
fd_errs.append(abs(acc - closed))
assert fd_errs[0] > fd_errs[2] or fd_errs[2] < 1e-8
closed2 = float(unit(torch.tensor([[z0]], dtype=torch.float64)).item())
assert closed2 == closed
assert fd_errs[-1] > fd_errs[2]python -m pytest tests/test_multipack_parity.py -k test_g2_collapse_no_closed_form_growth -qFields & methods
act_specfield
act_spec: ActivationSpec[Tensor]Declared ActivationSpec[Tensor] field on MultiPackUnit. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.
View source ↗meansfield
means: TensorDeclared Tensor field on MultiPackUnit. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.
View source ↗weightsfield
weights: TensorDeclared Tensor field on MultiPackUnit. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.
View source ↗__init__method
__init__(self, num_channels: int, spec: MultiPackSpec, *, base: str | ActivationSpec[Tensor]='sigmoid', learnable_means: bool=True, learnable_weights: bool=True, share_means: bool=True, dtype: torch.dtype | 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. |
specpositional-or-keyword | MultiPackSpecRequired | See the operation contract and implementation below. |
basekeyword-only | str | ActivationSpec[Tensor]Default: 'sigmoid' | See the operation contract and implementation below. |
learnable_meanskeyword-only | boolDefault: True | See the operation contract and implementation below. |
learnable_weightskeyword-only | boolDefault: True | See the operation contract and implementation below. |
share_meanskeyword-only | boolDefault: True | 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, raisesValueError(f'num_channels must be >= 1, got {num_channels}'). - When
self.act_spec.fastpath is None, raisesNotImplementedError(f'Activation {self.act_spec.name!r} has no closed-form derivative kernel').
Method implementation
def __init__(
self,
num_channels: int,
spec: MultiPackSpec,
*,
base: str | ActivationSpec[Tensor] = "sigmoid",
learnable_means: bool = True,
learnable_weights: bool = True,
share_means: bool = True,
dtype: torch.dtype | None = None,
) -> None:
super().__init__()
if num_channels < 1:
raise ValueError(f"num_channels must be >= 1, got {num_channels}")
self.num_channels = int(num_channels)
self.pack_spec = spec
self.orders: tuple[int, ...] = tuple(int(p.order) for p in spec.packs)
self.share_means = bool(share_means)
self.act_spec = base if isinstance(base, ActivationSpec) else get_activation(base)
if self.act_spec.fastpath is None:
raise NotImplementedError(
f"Activation {self.act_spec.name!r} has no closed-form derivative kernel"
)
probe = torch.zeros(1, dtype=torch.get_default_dtype())
for n in self.orders:
try:
self.act_spec.fastpath(probe, n)
except NotImplementedError:
raise
except Exception as exc: # pragma: no cover
raise NotImplementedError(
f"Activation {self.act_spec.name!r} cannot reach order {n}"
) from exc
dt = torch.get_default_dtype() if dtype is None else dtype
if share_means:
distinct = spec.distinct_means
mean_to_slot = {m: i for i, m in enumerate(distinct)}
self.mean_index: tuple[int, ...] = tuple(
mean_to_slot[p.mean] for p in spec.packs
)
means0 = torch.tensor(list(distinct), dtype=dt)
else:
self.mean_index = tuple(range(len(spec.packs)))
means0 = torch.tensor([p.mean for p in spec.packs], dtype=dt)
weights0 = torch.tensor([p.weight for p in spec.packs], dtype=dt)
if learnable_means:
self.means = nn.Parameter(means0)
else:
self.register_buffer("means", means0)
if learnable_weights:
self.weights = nn.Parameter(weights0)
else:
self.register_buffer("weights", weights0)forwardmethod
forward(self, z: Tensor) -> Tensorz shape (..., num_channels) -> same shape.
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, raisesValueError(f'expected z[..., {self.num_channels}], got shape {tuple(z.shape)}').
Method implementation
def forward(self, z: Tensor) -> Tensor:
"""``z`` shape ``(..., num_channels)`` -> same shape."""
if z.shape[-1] != self.num_channels:
raise ValueError(
f"expected z[..., {self.num_channels}], got shape {tuple(z.shape)}"
)
return multipack_response(
z,
self.means,
self.weights,
self.orders,
self.act_spec,
mean_index=self.mean_index,
)num_channelsattribute
num_channels = int(num_channels)Number of independent channels; z has shape (..., num_channels).
pack_specattribute
pack_spec: MultiPackSpec = specPublic instance attribute assigned by __init__. The source assignment is pack_spec = spec; it describes instance state, not a literal constructor default.
ordersattribute
orders: tuple[int, ...] = tuple((int(p.order) for p in spec.packs))Public instance attribute assigned by __init__. The source assignment is orders = tuple((int(p.order) for p in spec.packs)); it describes instance state, not a literal constructor default.
mean_indexattribute
mean_index: tuple[int, ...] = tuple((mean_to_slot[p.mean] for p in spec.packs))Public instance attribute assigned by __init__. The source assignment is mean_index = tuple((mean_to_slot[p.mean] for p in spec.packs)); it describes instance state, not a literal constructor default.
Read the implementation
class MultiPackUnit(nn.Module):
"""Channel-wise multi-pack Birkhoff unit.
Parameters
----------
num_channels:
Number of independent channels; ``z`` has shape ``(..., num_channels)``.
spec:
:class:`~omnibias.core.multipack.MultiPackSpec` describing orders,
initial means, and initial outer weights.
base:
Activation name or :class:`~omnibias.torch.activations.ActivationSpec`.
learnable_means / learnable_weights:
Whether means / outer weights are ``nn.Parameter`` or frozen buffers.
share_means:
If true (default), packs with equal means in ``spec`` share one mean
Parameter (one base evaluation per mean with a registered tower).
dtype:
Parameter dtype; ``None`` resolves to ``torch.get_default_dtype()``.
"""
act_spec: ActivationSpec[Tensor]
means: Tensor
weights: Tensor
def __init__(
self,
num_channels: int,
spec: MultiPackSpec,
*,
base: str | ActivationSpec[Tensor] = "sigmoid",
learnable_means: bool = True,
learnable_weights: bool = True,
share_means: bool = True,
dtype: torch.dtype | None = None,
) -> None:
super().__init__()
if num_channels < 1:
raise ValueError(f"num_channels must be >= 1, got {num_channels}")
self.num_channels = int(num_channels)
self.pack_spec = spec
self.orders: tuple[int, ...] = tuple(int(p.order) for p in spec.packs)
self.share_means = bool(share_means)
self.act_spec = base if isinstance(base, ActivationSpec) else get_activation(base)
if self.act_spec.fastpath is None:
raise NotImplementedError(
f"Activation {self.act_spec.name!r} has no closed-form derivative kernel"
)
probe = torch.zeros(1, dtype=torch.get_default_dtype())
for n in self.orders:
try:
self.act_spec.fastpath(probe, n)
except NotImplementedError:
raise
except Exception as exc: # pragma: no cover
raise NotImplementedError(
f"Activation {self.act_spec.name!r} cannot reach order {n}"
) from exc
dt = torch.get_default_dtype() if dtype is None else dtype
if share_means:
distinct = spec.distinct_means
mean_to_slot = {m: i for i, m in enumerate(distinct)}
self.mean_index: tuple[int, ...] = tuple(
mean_to_slot[p.mean] for p in spec.packs
)
means0 = torch.tensor(list(distinct), dtype=dt)
else:
self.mean_index = tuple(range(len(spec.packs)))
means0 = torch.tensor([p.mean for p in spec.packs], dtype=dt)
weights0 = torch.tensor([p.weight for p in spec.packs], dtype=dt)
if learnable_means:
self.means = nn.Parameter(means0)
else:
self.register_buffer("means", means0)
if learnable_weights:
self.weights = nn.Parameter(weights0)
else:
self.register_buffer("weights", weights0)
def forward(self, z: Tensor) -> Tensor:
"""``z`` shape ``(..., num_channels)`` -> same shape."""
if z.shape[-1] != self.num_channels:
raise ValueError(
f"expected z[..., {self.num_channels}], got shape {tuple(z.shape)}"
)
return multipack_response(
z,
self.means,
self.weights,
self.orders,
self.act_spec,
mean_index=self.mean_index,
)Other public import paths
omnibias.torch.BirkhoffOMBUomnibias.torch.MultiPackUnitomnibias.torch.multipack.BirkhoffOMBUomnibias.torch.multipack.MultiPackUnit
multipack_response
omnibias.torch.multipack_responseSource ↗multipack_response(z: Tensor, means: Tensor, weights: Tensor, orders: tuple[int, ...], spec: ActivationSpec[Tensor], *, mean_index: tuple[int, ...] | None=None) -> TensorEvaluate sum_g c_g sigma^(n_g)(z + mu_g) with a closed-form tower.
orders must be a Python tuple of static ints. When mean_index is
given, pack g reads means[mean_index[g]] (shared-mean layout);
otherwise means has length G and pack g reads means[g].
Parameters
| Name | Type & default | Contract |
|---|---|---|
zpositional-or-keyword | TensorRequired | See the operation contract and implementation below. |
meanspositional-or-keyword | TensorRequired | See the operation contract and implementation below. |
weightspositional-or-keyword | TensorRequired | See the operation contract and implementation below. |
orderspositional-or-keyword | tuple[int, ...]Required | See the operation contract and implementation below. |
specpositional-or-keyword | ActivationSpec[Tensor]Required | See the operation contract and implementation below. |
mean_indexkeyword-only | tuple[int, ...] | NoneDefault: None | See the operation contract and implementation below. |
TensorExplicit validation & failure cases
- When
spec.fastpath is None, raisesNotImplementedError(f'Activation {spec.name!r} has no closed-form derivative kernel'). - When
int(weights.shape[-1]) != G, raisesValueError('weights length must equal number of packs'). - When
int(means.shape[-1]) != G, raisesValueError('means length must equal number of packs'). - When
len(mean_index) != G, raisesValueError('mean_index length must equal number of packs'). - When
n < 0, raisesValueError(f'order must be >= 0, got {n}').
These are explicit source guards. Backend operations may also raise for invalid shapes, dtypes, or unsupported devices.
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.
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" -qRead the implementation
def multipack_response(
z: Tensor,
means: Tensor,
weights: Tensor,
orders: tuple[int, ...],
spec: ActivationSpec[Tensor],
*,
mean_index: tuple[int, ...] | None = None,
) -> Tensor:
"""Evaluate ``sum_g c_g sigma^(n_g)(z + mu_g)`` with a closed-form tower.
``orders`` must be a Python tuple of static ints. When ``mean_index`` is
given, pack ``g`` reads ``means[mean_index[g]]`` (shared-mean layout);
otherwise ``means`` has length ``G`` and pack ``g`` reads ``means[g]``.
"""
if spec.fastpath is None:
raise NotImplementedError(
f"Activation {spec.name!r} has no closed-form derivative kernel"
)
G = len(orders)
if int(weights.shape[-1]) != G:
raise ValueError("weights length must equal number of packs")
if mean_index is None:
if int(means.shape[-1]) != G:
raise ValueError("means length must equal number of packs")
index = tuple(range(G))
else:
if len(mean_index) != G:
raise ValueError("mean_index length must equal number of packs")
index = mean_index
for n in orders:
if n < 0:
raise ValueError(f"order must be >= 0, got {n}")
fp = spec.fastpath
slots = sorted(set(index))
slot_u: dict[int, Tensor] = {s: z + means[s] for s in slots}
towers: dict[int, Tensor] = {}
if spec.tower is not None:
for slot in slots:
highest = max(n for g, n in enumerate(orders) if index[g] == slot)
towers[slot] = spec.tower(slot_u[slot], highest)
out: Tensor | None = None
for g, n in enumerate(orders):
value = towers[index[g]][n] if spec.tower is not None else fp(slot_u[index[g]], n)
term = weights[g] * value
out = term if out is None else out + term
assert out is not None
return outOther public import paths
omnibias.torch.multipack_responseomnibias.torch.multipack.multipack_response