omnibias.partition.arrangement.torch
Trainable soft partitions of unity, regional experts, temperature hardening, and scoped membership bounds.
import omnibias.partition.arrangement.torchModule contract & mathematical conventions
Soft arrangement membership (torch; theory 01-03).
beta -> inf is temperature collapse, not founding delta -> 0.
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.
margin
omnibias.partition.arrangement.torch.marginSource ↗margin(arr: Arrangement, x: Tensor) -> TensorThe module contract describes the context for this operation. Its signature, explicit validation, return expressions, and implementation are recorded below.
Parameters
| Name | Type & default | Contract |
|---|---|---|
arrpositional-or-keyword | ArrangementRequired | See the operation contract and implementation below. |
xpositional-or-keyword | TensorRequired | See the operation contract and implementation below. |
TensorRead the implementation
def margin(arr: Arrangement, x: Tensor) -> Tensor:
w = torch.as_tensor(arr.normals, dtype=x.dtype, device=x.device)
t = torch.as_tensor(arr.offsets, dtype=x.dtype, device=x.device)
xv = x if x.ndim == 2 else x.unsqueeze(0)
z = xv @ w.T - t
return z.abs().amin(dim=-1)soft_membership
omnibias.partition.arrangement.torch.soft_membershipSource ↗soft_membership(arr: Arrangement, x: Tensor, signs: Sequence[int], *, beta: float) -> TensorThe module contract describes the context for this operation. Its signature, explicit validation, return expressions, and implementation are recorded below.
Parameters
| Name | Type & default | Contract |
|---|---|---|
arrpositional-or-keyword | ArrangementRequired | See the operation contract and implementation below. |
xpositional-or-keyword | TensorRequired | See the operation contract and implementation below. |
signspositional-or-keyword | Sequence[int]Required | See the operation contract and implementation below. |
betakeyword-only | floatRequired | See the operation contract and implementation below. |
TensorExplicit validation & failure cases
- When
beta <= 0.0, raisesValueError('beta must be > 0 (temperature collapse axis)').
These are explicit source guards. Backend operations may also raise for invalid shapes, dtypes, or unsupported devices.
Executed regression · g4 torch jax parity
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_g4_torch_jax_parity() -> None:
pytest.importorskip("jax")
import jax
import jax.numpy as jnp
import torch
from omnibias.partition.arrangement.jax import soft_membership as sm_jax
from omnibias.partition.arrangement.torch import soft_membership as sm_torch
jax.config.update("jax_enable_x64", True)
rng = np.random.default_rng(2)
arr = general_position_normals(4, 2, rng)
x = rng.normal(size=(6, 2))
signs = (1, -1, 1, -1)
t = sm_torch(arr, torch.as_tensor(x, dtype=torch.float64), signs, beta=2.5)
j = sm_jax(arr, jnp.asarray(x, dtype=jnp.float64), signs, beta=2.5)
np.testing.assert_allclose(t.detach().cpu().numpy(), np.asarray(j), rtol=0.0, atol=1e-14)python -m pytest "packages/omnibias-partition/tests/test_arrangement.py::test_g4_torch_jax_parity" -qRead the implementation
def soft_membership(
arr: Arrangement,
x: Tensor,
signs: Sequence[int],
*,
beta: float,
) -> Tensor:
if beta <= 0.0:
raise ValueError("beta must be > 0 (temperature collapse axis)")
w = torch.as_tensor(arr.normals, dtype=x.dtype, device=x.device)
t = torch.as_tensor(arr.offsets, dtype=x.dtype, device=x.device)
xv = x if x.ndim == 2 else x.unsqueeze(0)
z = xv @ w.T - t
s = torch.as_tensor(list(signs), dtype=x.dtype, device=x.device)
return torch.sigmoid(beta * s * z).prod(dim=-1)