omnibias.torch.architectures.piratenet
PyTorch activation derivatives, trainable operator layers, network jets, and direct differential-operator contractions.
import omnibias.torch.architectures.piratenetModule contract & mathematical conventions
Identity-initialized residual networks for physics-informed models.
The zero skip parameter selects the embedding; Fourier features compose with the shared field architecture.
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.
PirateNet
omnibias.torch.architectures.PirateNetSource ↗class PirateNet(nn.Module)Trainable PirateNet wrapping the functional apply.
Parameters
| Name | Type & default | Contract |
|---|---|---|
cfgpositional-or-keyword | PirateNetConfig | NoneDefault: None | See the operation contract and implementation below. |
Executed regression · module 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_module_forward_shape() -> None:
net = PirateNet(PirateNetConfig(in_dim=1, hidden=5, n_layers=1, seed=0))
y = net(torch.linspace(-1.0, 1.0, 7, dtype=torch.float64)[:, None])
assert y.shape == (7,)python -m pytest "packages/omnibias-torch/tests/test_piratenet.py::test_module_forward_shape" -qFields & methods
__init__method
__init__(self, cfg: PirateNetConfig | None=None) -> NoneSee the class contract above and the exact method implementation below.
Parameters
| Name | Type & default | Contract |
|---|---|---|
cfgpositional-or-keyword | PirateNetConfig | NoneDefault: None | See the operation contract and implementation below. |
Returns: None
Executed regression · module 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_module_forward_shape() -> None:
net = PirateNet(PirateNetConfig(in_dim=1, hidden=5, n_layers=1, seed=0))
y = net(torch.linspace(-1.0, 1.0, 7, dtype=torch.float64)[:, None])
assert y.shape == (7,)python -m pytest "packages/omnibias-torch/tests/test_piratenet.py::test_module_forward_shape" -qMethod implementation
def __init__(self, cfg: PirateNetConfig | None = None) -> None:
super().__init__()
cfg = PirateNetConfig() if cfg is None else cfg
raw = init_pirate_params(cfg)
self.in_dim = int(cfg.in_dim)
self.hidden = int(cfg.hidden)
self.n_layers = int(cfg.n_layers)
self.out_dim = int(cfg.out_dim)
self.We = nn.Parameter(raw["We"])
self.be = nn.Parameter(raw["be"])
self.Wu = nn.Parameter(raw["Wu"])
self.bu = nn.Parameter(raw["bu"])
self.Wv = nn.Parameter(raw["Wv"])
self.bv = nn.Parameter(raw["bv"])
self.Wout = nn.Parameter(raw["Wout"])
self.bout = nn.Parameter(raw["bout"])
self.alpha = nn.Parameter(raw["alpha"])
blocks = []
for block in raw["blocks"]:
blocks.append(
nn.ParameterDict(
{
"W1": nn.Parameter(block["W1"]),
"b1": nn.Parameter(block["b1"]),
"W2": nn.Parameter(block["W2"]),
"b2": nn.Parameter(block["b2"]),
"W3": nn.Parameter(block["W3"]),
"b3": nn.Parameter(block["b3"]),
}
)
)
self.blocks = nn.ModuleList(blocks)forwardmethod
forward(self, x: Tensor) -> TensorSee the class contract above and the exact method implementation below.
Parameters
| Name | Type & default | Contract |
|---|---|---|
xpositional-or-keyword | TensorRequired | See the operation contract and implementation below. |
Returns: Tensor
Executed regression · module 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_module_forward_shape() -> None:
net = PirateNet(PirateNetConfig(in_dim=1, hidden=5, n_layers=1, seed=0))
y = net(torch.linspace(-1.0, 1.0, 7, dtype=torch.float64)[:, None])
assert y.shape == (7,)python -m pytest "packages/omnibias-torch/tests/test_piratenet.py::test_module_forward_shape" -qMethod implementation
def forward(self, x: Tensor) -> Tensor:
return pirate_apply(self._params(), x)in_dimattribute
in_dim = int(cfg.in_dim)Public instance attribute assigned by __init__. The source assignment is in_dim = int(cfg.in_dim); it describes instance state, not a literal constructor default.
n_layersattribute
n_layers = int(cfg.n_layers)Public instance attribute assigned by __init__. The source assignment is n_layers = int(cfg.n_layers); it describes instance state, not a literal constructor default.
out_dimattribute
out_dim = int(cfg.out_dim)Public instance attribute assigned by __init__. The source assignment is out_dim = int(cfg.out_dim); it describes instance state, not a literal constructor default.
Weattribute
We = nn.Parameter(raw['We'])Public instance attribute assigned by __init__. The source assignment is We = nn.Parameter(raw['We']); it describes instance state, not a literal constructor default.
beattribute
be = nn.Parameter(raw['be'])Public instance attribute assigned by __init__. The source assignment is be = nn.Parameter(raw['be']); it describes instance state, not a literal constructor default.
Wuattribute
Wu = nn.Parameter(raw['Wu'])Public instance attribute assigned by __init__. The source assignment is Wu = nn.Parameter(raw['Wu']); it describes instance state, not a literal constructor default.
buattribute
bu = nn.Parameter(raw['bu'])Public instance attribute assigned by __init__. The source assignment is bu = nn.Parameter(raw['bu']); it describes instance state, not a literal constructor default.
Wvattribute
Wv = nn.Parameter(raw['Wv'])Public instance attribute assigned by __init__. The source assignment is Wv = nn.Parameter(raw['Wv']); it describes instance state, not a literal constructor default.
bvattribute
bv = nn.Parameter(raw['bv'])Public instance attribute assigned by __init__. The source assignment is bv = nn.Parameter(raw['bv']); it describes instance state, not a literal constructor default.
Woutattribute
Wout = nn.Parameter(raw['Wout'])Public instance attribute assigned by __init__. The source assignment is Wout = nn.Parameter(raw['Wout']); it describes instance state, not a literal constructor default.
boutattribute
bout = nn.Parameter(raw['bout'])Public instance attribute assigned by __init__. The source assignment is bout = nn.Parameter(raw['bout']); it describes instance state, not a literal constructor default.
alphaattribute
alpha = nn.Parameter(raw['alpha'])Public instance attribute assigned by __init__. The source assignment is alpha = nn.Parameter(raw['alpha']); it describes instance state, not a literal constructor default.
blocksattribute
blocks = nn.ModuleList(blocks)Public instance attribute assigned by __init__. The source assignment is blocks = nn.ModuleList(blocks); it describes instance state, not a literal constructor default.
Read the implementation
class PirateNet(nn.Module):
"""Trainable PirateNet wrapping the functional apply."""
def __init__(self, cfg: PirateNetConfig | None = None) -> None:
super().__init__()
cfg = PirateNetConfig() if cfg is None else cfg
raw = init_pirate_params(cfg)
self.in_dim = int(cfg.in_dim)
self.hidden = int(cfg.hidden)
self.n_layers = int(cfg.n_layers)
self.out_dim = int(cfg.out_dim)
self.We = nn.Parameter(raw["We"])
self.be = nn.Parameter(raw["be"])
self.Wu = nn.Parameter(raw["Wu"])
self.bu = nn.Parameter(raw["bu"])
self.Wv = nn.Parameter(raw["Wv"])
self.bv = nn.Parameter(raw["bv"])
self.Wout = nn.Parameter(raw["Wout"])
self.bout = nn.Parameter(raw["bout"])
self.alpha = nn.Parameter(raw["alpha"])
blocks = []
for block in raw["blocks"]:
blocks.append(
nn.ParameterDict(
{
"W1": nn.Parameter(block["W1"]),
"b1": nn.Parameter(block["b1"]),
"W2": nn.Parameter(block["W2"]),
"b2": nn.Parameter(block["b2"]),
"W3": nn.Parameter(block["W3"]),
"b3": nn.Parameter(block["b3"]),
}
)
)
self.blocks = nn.ModuleList(blocks)
def _params(self) -> dict[str, Any]:
return {
"We": self.We,
"be": self.be,
"Wu": self.Wu,
"bu": self.bu,
"Wv": self.Wv,
"bv": self.bv,
"Wout": self.Wout,
"bout": self.bout,
"alpha": self.alpha,
"blocks": [
{
k: cast(nn.ParameterDict, block)[k]
for k in ("W1", "b1", "W2", "b2", "W3", "b3")
}
for block in self.blocks
],
}
def forward(self, x: Tensor) -> Tensor:
return pirate_apply(self._params(), x)Other public import paths
omnibias.torch.architectures.PirateNetomnibias.torch.architectures.piratenet.PirateNet
PirateNetConfig
omnibias.torch.architectures.PirateNetConfigSource ↗class PirateNetConfigDataclass constructor fields; inherited fields, if any, are declared by the base class. A field(...) expression describes a default factory or field options, not a literal argument default.
PirateNet width / depth. α starts at 0.
Parameters
| Name | Type & default | Contract |
|---|---|---|
in_dimdataclass field | intDefault: 1 | See the operation contract and implementation below. |
hiddendataclass field | intDefault: 24 | See the operation contract and implementation below. |
n_layersdataclass field | intDefault: 2 | See the operation contract and implementation below. |
out_dimdataclass field | intDefault: 1 | See the operation contract and implementation below. |
seeddataclass field | intDefault: 0 | See the operation contract and implementation below. |
Regression example · module forward shape
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_module_forward_shape() -> None:
net = PirateNet(PirateNetConfig(in_dim=1, hidden=5, n_layers=1, seed=0))
y = net(torch.linspace(-1.0, 1.0, 7, dtype=torch.float64)[:, None])
assert y.shape == (7,)python -m pytest packages/omnibias-torch/tests/test_piratenet.py -k test_module_forward_shape -qRegression example · alpha zero is embedding
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_alpha_zero_is_embedding() -> None:
cfg = PirateNetConfig(in_dim=2, hidden=6, n_layers=3, seed=1)
params = init_pirate_params(cfg)
assert float(torch.max(torch.abs(params["alpha"]))) == 0.0
x = torch.linspace(-1.0, 1.0, 10, dtype=torch.float64)
coords = torch.stack([x, x * x], dim=-1)
feat = pirate_features(params, coords)
emb = torch.tanh(coords @ params["We"].T + params["be"])
assert float(torch.max(torch.abs(feat - emb))) < 1e-14python -m pytest packages/omnibias-torch/tests/test_piratenet.py -k test_alpha_zero_is_embedding -qFields & methods
in_dimfield
in_dim: int = 1Declared int field on PirateNetConfig. Default expression: 1.
n_layersfield
n_layers: int = 2Declared int field on PirateNetConfig. Default expression: 2.
out_dimfield
out_dim: int = 1Declared int field on PirateNetConfig. Default expression: 1.
Read the implementation
class PirateNetConfig:
"""PirateNet width / depth. ``α`` starts at 0."""
in_dim: int = 1
hidden: int = 24
n_layers: int = 2
out_dim: int = 1
seed: int = 0Other public import paths
omnibias.torch.architectures.PirateNetConfigomnibias.torch.architectures.piratenet.PirateNetConfig
init_pirate_params
omnibias.torch.architectures.init_pirate_paramsSource ↗init_pirate_params(cfg: PirateNetConfig) -> dict[str, Any]Identity-skip PirateNet (alpha=0) plus a zero readout.
Parameters
| Name | Type & default | Contract |
|---|---|---|
cfgpositional-or-keyword | PirateNetConfigRequired | See the operation contract and implementation below. |
dict[str, Any]Executed regression · module 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_module_forward_shape() -> None:
net = PirateNet(PirateNetConfig(in_dim=1, hidden=5, n_layers=1, seed=0))
y = net(torch.linspace(-1.0, 1.0, 7, dtype=torch.float64)[:, None])
assert y.shape == (7,)python -m pytest "packages/omnibias-torch/tests/test_piratenet.py::test_module_forward_shape" -qRegression example · alpha zero is embedding
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_alpha_zero_is_embedding() -> None:
cfg = PirateNetConfig(in_dim=2, hidden=6, n_layers=3, seed=1)
params = init_pirate_params(cfg)
assert float(torch.max(torch.abs(params["alpha"]))) == 0.0
x = torch.linspace(-1.0, 1.0, 10, dtype=torch.float64)
coords = torch.stack([x, x * x], dim=-1)
feat = pirate_features(params, coords)
emb = torch.tanh(coords @ params["We"].T + params["be"])
assert float(torch.max(torch.abs(feat - emb))) < 1e-14python -m pytest packages/omnibias-torch/tests/test_piratenet.py -k test_alpha_zero_is_embedding -qRead the implementation
def init_pirate_params(cfg: PirateNetConfig) -> dict[str, Any]:
"""Identity-skip PirateNet (``alpha=0``) plus a zero readout."""
hid = int(cfg.hidden)
n_layers = int(cfg.n_layers)
in_dim = int(cfg.in_dim)
g = torch.Generator().manual_seed(int(cfg.seed))
if int(cfg.out_dim) == 1:
wout: Tensor = torch.zeros(hid, dtype=torch.float64)
bout: Tensor = torch.zeros((), dtype=torch.float64)
else:
wout = torch.zeros(hid, int(cfg.out_dim), dtype=torch.float64)
bout = torch.zeros(int(cfg.out_dim), dtype=torch.float64)
params: dict[str, Any] = {
"We": _glorot((hid, in_dim), g),
"be": torch.zeros(hid, dtype=torch.float64),
"Wu": _glorot((hid, hid), g),
"bu": torch.zeros(hid, dtype=torch.float64),
"Wv": _glorot((hid, hid), g),
"bv": torch.zeros(hid, dtype=torch.float64),
"Wout": wout,
"bout": bout,
"alpha": torch.zeros(n_layers, dtype=torch.float64),
"blocks": [],
}
for _ in range(n_layers):
params["blocks"].append(
{
"W1": _glorot((hid, hid), g),
"b1": torch.zeros(hid, dtype=torch.float64),
"W2": _glorot((hid, hid), g),
"b2": torch.zeros(hid, dtype=torch.float64),
"W3": _glorot((hid, hid), g),
"b3": torch.zeros(hid, dtype=torch.float64),
}
)
params["blocks"] = tuple(params["blocks"])
return paramsOther public import paths
omnibias.torch.architectures.init_pirate_paramsomnibias.torch.architectures.piratenet.init_pirate_params
pirate_apply
omnibias.torch.architectures.pirate_applySource ↗pirate_apply(params: dict[str, Any], x: Tensor) -> TensorLinear readout of pirate_features.
Parameters
| Name | Type & default | Contract |
|---|---|---|
paramspositional-or-keyword | dict[str, Any]Required | See the operation contract and implementation below. |
xpositional-or-keyword | TensorRequired | See the operation contract and implementation below. |
TensorExecuted regression · module 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_module_forward_shape() -> None:
net = PirateNet(PirateNetConfig(in_dim=1, hidden=5, n_layers=1, seed=0))
y = net(torch.linspace(-1.0, 1.0, 7, dtype=torch.float64)[:, None])
assert y.shape == (7,)python -m pytest "packages/omnibias-torch/tests/test_piratenet.py::test_module_forward_shape" -qRegression example · pirate apply parity shared weights
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_pirate_apply_parity_shared_weights() -> None:
cfg = JaxCfg(in_dim=2, hidden=5, n_layers=2, seed=3)
jparams = jax_init(cfg)
tparams = _to_torch(jparams)
x = np.array([[0.1, 0.2], [-0.3, 0.4], [0.0, 0.0]], dtype=np.float64)
yj = np.asarray(jax_apply(jparams, jnp.asarray(x)))
yt = torch_apply(tparams, torch.tensor(x)).detach().numpy()
np.testing.assert_allclose(yj, yt, atol=1e-14, rtol=0.0)python -m pytest tests/test_piratenet_parity.py -k test_pirate_apply_parity_shared_weights -qRead the implementation
def pirate_apply(params: dict[str, Any], x: Tensor) -> Tensor:
"""Linear readout of :func:`pirate_features`."""
feat = pirate_features(params, x)
return torch.as_tensor(feat @ params["Wout"] + params["bout"])Other public import paths
omnibias.torch.architectures.pirate_applyomnibias.torch.architectures.piratenet.pirate_apply
pirate_features
omnibias.torch.architectures.pirate_featuresSource ↗pirate_features(params: dict[str, Any], x: Tensor) -> TensorPenultimate features, shape (..., hidden). alpha=0 is identity.
Parameters
| Name | Type & default | Contract |
|---|---|---|
paramspositional-or-keyword | dict[str, Any]Required | See the operation contract and implementation below. |
xpositional-or-keyword | TensorRequired | See the operation contract and implementation below. |
TensorExecuted regression · module 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_module_forward_shape() -> None:
net = PirateNet(PirateNetConfig(in_dim=1, hidden=5, n_layers=1, seed=0))
y = net(torch.linspace(-1.0, 1.0, 7, dtype=torch.float64)[:, None])
assert y.shape == (7,)python -m pytest "packages/omnibias-torch/tests/test_piratenet.py::test_module_forward_shape" -qRegression example · alpha zero is embedding
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_alpha_zero_is_embedding() -> None:
cfg = PirateNetConfig(in_dim=2, hidden=6, n_layers=3, seed=1)
params = init_pirate_params(cfg)
assert float(torch.max(torch.abs(params["alpha"]))) == 0.0
x = torch.linspace(-1.0, 1.0, 10, dtype=torch.float64)
coords = torch.stack([x, x * x], dim=-1)
feat = pirate_features(params, coords)
emb = torch.tanh(coords @ params["We"].T + params["be"])
assert float(torch.max(torch.abs(feat - emb))) < 1e-14python -m pytest packages/omnibias-torch/tests/test_piratenet.py -k test_alpha_zero_is_embedding -qRead the implementation
def pirate_features(params: dict[str, Any], x: Tensor) -> Tensor:
"""Penultimate features, shape ``(..., hidden)``. ``alpha=0`` is identity."""
coords = torch.as_tensor(x, dtype=torch.float64)
if coords.ndim == 1:
coords = coords[None, :]
squeeze = True
else:
squeeze = False
h = torch.tanh(coords @ params["We"].T + params["be"])
u = torch.tanh(h @ params["Wu"].T + params["bu"])
v = torch.tanh(h @ params["Wv"].T + params["bv"])
for i, block in enumerate(params["blocks"]):
identity = h
z = torch.tanh(h @ block["W1"].T + block["b1"])
z = z * u + (1.0 - z) * v
z = torch.tanh(z @ block["W2"].T + block["b2"])
z = z * u + (1.0 - z) * v
z = torch.tanh(z @ block["W3"].T + block["b3"])
alpha = params["alpha"][i]
h = alpha * z + (1.0 - alpha) * identity
if squeeze:
return torch.as_tensor(h[0])
return torch.as_tensor(h)Other public import paths
omnibias.torch.architectures.pirate_featuresomnibias.torch.architectures.piratenet.pirate_features