omnibias.jax.architectures.piratenet
JAX activation derivatives, functional network jets, direct contractions, and parameter-differentiable numerical models.
import omnibias.jax.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.
PirateNetConfig
omnibias.jax.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 · 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 -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(jnp.max(jnp.abs(params["alpha"]))) == 0.0
x = jnp.linspace(-1.0, 1.0, 10, dtype=jnp.float64)
coords = jnp.stack([x, x * x], axis=-1)
feat = pirate_features(params, coords)
emb = jnp.tanh(coords @ params["We"].T + params["be"])
assert float(jnp.max(jnp.abs(feat - emb))) < 1e-14python -m pytest packages/omnibias-jax/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.jax.architectures.PirateNetConfigomnibias.jax.architectures.piratenet.PirateNetConfig
init_pirate_params
omnibias.jax.architectures.init_pirate_paramsSource ↗init_pirate_params(cfg: PirateNetConfig, *, key: Array | None=None) -> 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. |
keykeyword-only | Array | NoneDefault: None | See the operation contract and implementation below. |
dict[str, Any]Executed regression · alpha zero is embedding
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_alpha_zero_is_embedding() -> None:
cfg = PirateNetConfig(in_dim=2, hidden=6, n_layers=3, seed=1)
params = init_pirate_params(cfg)
assert float(jnp.max(jnp.abs(params["alpha"]))) == 0.0
x = jnp.linspace(-1.0, 1.0, 10, dtype=jnp.float64)
coords = jnp.stack([x, x * x], axis=-1)
feat = pirate_features(params, coords)
emb = jnp.tanh(coords @ params["We"].T + params["be"])
assert float(jnp.max(jnp.abs(feat - emb))) < 1e-14python -m pytest "packages/omnibias-jax/tests/test_piratenet.py::test_alpha_zero_is_embedding" -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 init_pirate_params(
cfg: PirateNetConfig,
*,
key: Array | None = None,
) -> 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)
rng = jax.random.PRNGKey(int(cfg.seed)) if key is None else key
keys = jax.random.split(rng, 6 + 3 * n_layers)
if int(cfg.out_dim) == 1:
wout: Array = jnp.zeros((hid,), dtype=jnp.float64)
bout: Array = jnp.zeros((), dtype=jnp.float64)
else:
wout = jnp.zeros((hid, int(cfg.out_dim)), dtype=jnp.float64)
bout = jnp.zeros((int(cfg.out_dim),), dtype=jnp.float64)
params: dict[str, Any] = {
"We": _glorot(keys[0], (hid, in_dim)),
"be": jnp.zeros((hid,), dtype=jnp.float64),
"Wu": _glorot(keys[1], (hid, hid)),
"bu": jnp.zeros((hid,), dtype=jnp.float64),
"Wv": _glorot(keys[2], (hid, hid)),
"bv": jnp.zeros((hid,), dtype=jnp.float64),
"Wout": wout,
"bout": bout,
"alpha": jnp.zeros((n_layers,), dtype=jnp.float64),
"blocks": [],
}
for i in range(n_layers):
k1, k2, k3 = keys[6 + 3 * i : 6 + 3 * i + 3]
params["blocks"].append(
{
"W1": _glorot(k1, (hid, hid)),
"b1": jnp.zeros((hid,), dtype=jnp.float64),
"W2": _glorot(k2, (hid, hid)),
"b2": jnp.zeros((hid,), dtype=jnp.float64),
"W3": _glorot(k3, (hid, hid)),
"b3": jnp.zeros((hid,), dtype=jnp.float64),
}
)
params["blocks"] = tuple(params["blocks"])
return paramsOther public import paths
omnibias.jax.architectures.init_pirate_paramsomnibias.jax.architectures.piratenet.init_pirate_params
pirate_apply
omnibias.jax.architectures.pirate_applySource ↗pirate_apply(params: dict[str, Any], x: Array) -> ArrayLinear 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 | ArrayRequired | See the operation contract and implementation below. |
ArrayExecuted regression · one d residual smoke decreases
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_one_d_residual_smoke_decreases() -> None:
cfg = PirateNetConfig(in_dim=1, hidden=8, n_layers=1, seed=2)
params = init_pirate_params(cfg)
x = jnp.linspace(-1.0, 1.0, 32, dtype=jnp.float64)[:, None]
target = jnp.sin(jnp.pi * x[:, 0])
def loss(th: dict) -> jnp.ndarray:
pred = pirate_apply(th, x)
return jnp.mean((pred - target) ** 2)
l0 = float(loss(params))
grad = jax.grad(loss)(params)
params = jax.tree_util.tree_map(lambda p, g: p - 0.05 * g, params, grad)
l1 = float(loss(params))
assert l1 < l0python -m pytest "packages/omnibias-jax/tests/test_piratenet.py::test_one_d_residual_smoke_decreases" -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: Array) -> Array:
"""Linear readout of :func:`pirate_features`."""
feat = pirate_features(params, x)
return jnp.asarray(feat @ params["Wout"] + params["bout"])Other public import paths
omnibias.jax.architectures.pirate_applyomnibias.jax.architectures.piratenet.pirate_apply
pirate_features
omnibias.jax.architectures.pirate_featuresSource ↗pirate_features(params: dict[str, Any], x: Array) -> ArrayPenultimate 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 | ArrayRequired | See the operation contract and implementation below. |
ArrayExecuted regression · alpha zero is embedding
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_alpha_zero_is_embedding() -> None:
cfg = PirateNetConfig(in_dim=2, hidden=6, n_layers=3, seed=1)
params = init_pirate_params(cfg)
assert float(jnp.max(jnp.abs(params["alpha"]))) == 0.0
x = jnp.linspace(-1.0, 1.0, 10, dtype=jnp.float64)
coords = jnp.stack([x, x * x], axis=-1)
feat = pirate_features(params, coords)
emb = jnp.tanh(coords @ params["We"].T + params["be"])
assert float(jnp.max(jnp.abs(feat - emb))) < 1e-14python -m pytest "packages/omnibias-jax/tests/test_piratenet.py::test_alpha_zero_is_embedding" -qRead the implementation
def pirate_features(params: dict[str, Any], x: Array) -> Array:
"""Penultimate features, shape ``(..., hidden)``. ``alpha=0`` is identity."""
coords = jnp.asarray(x, dtype=jnp.float64)
if coords.ndim == 1:
coords = coords[None, :]
squeeze = True
else:
squeeze = False
h = jnp.tanh(coords @ params["We"].T + params["be"])
u = jnp.tanh(h @ params["Wu"].T + params["bu"])
v = jnp.tanh(h @ params["Wv"].T + params["bv"])
for i, block in enumerate(params["blocks"]):
identity = h
z = jnp.tanh(h @ block["W1"].T + block["b1"])
z = z * u + (1.0 - z) * v
z = jnp.tanh(z @ block["W2"].T + block["b2"])
z = z * u + (1.0 - z) * v
z = jnp.tanh(z @ block["W3"].T + block["b3"])
alpha = params["alpha"][i]
h = alpha * z + (1.0 - alpha) * identity
if squeeze:
return jnp.asarray(h[0])
return jnp.asarray(h)Other public import paths
omnibias.jax.architectures.pirate_featuresomnibias.jax.architectures.piratenet.pirate_features