Skip to main content

PyTorch quickstart

A minimal, runnable example. Install with pip install omnibias-torch.

import torch
from omnibias.torch import OMBU, OperatorBlock, cmbLinear

# 1. OMBU: a trainable K-bias operator, drop-in for an activation.
ombu = OMBU(num_channels=4, K=2, base="tanh")
out = ombu(torch.zeros(8, 4))
print("OMBU out:", out.shape)

# 2. OperatorBlock: a typed scalar operator.
# op ∈ {identity, grad, laplacian, derivative, band, integral}
laplacian = OperatorBlock(channels=8, op="laplacian", base="gaussian")

# 3. cmbLinear: nn.Linear with an inline OperatorBlock.
fc = cmbLinear(in_features=128, out_features=64, op="identity", base="tanh")

# Compose into a normal module.
model = torch.nn.Sequential(
cmbLinear(16, 64, op="identity", base="tanh"),
cmbLinear(64, 1, op="identity", base="tanh"),
)
y = model(torch.randn(32, 16))
print("model out:", y.shape)

The derivative tower directly

from omnibias.torch import get_activation

spec = get_activation("tanh")
z = torch.linspace(-2, 2, 5)
for n in range(4):
print(n, spec.fastpath(z, n)) # σ, σ', σ'', σ''' - closed form

Next