Skip to main content

omnibias.partition.arrangement.torch

0.1.0a2AlphaPython >=3.10AGPL / commercial

Trainable soft partitions of unity, regional experts, temperature hardening, and scoped membership bounds.

import omnibias.partition.arrangement.torch
Module 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.

FUNCTION

margin​

omnibias.partition.arrangement.torch.marginSource ↗
margin(arr: Arrangement, x: Tensor) -> Tensor

The module contract describes the context for this operation. Its signature, explicit validation, return expressions, and implementation are recorded below.

Parameters

NameType & defaultContract
arrpositional-or-keywordArrangementRequiredSee the operation contract and implementation below.
xpositional-or-keywordTensorRequiredSee the operation contract and implementation below.
ReturnsTensor
Read 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)
FUNCTION

soft_membership​

omnibias.partition.arrangement.torch.soft_membershipSource ↗
soft_membership(arr: Arrangement, x: Tensor, signs: Sequence[int], *, beta: float) -> Tensor

The module contract describes the context for this operation. Its signature, explicit validation, return expressions, and implementation are recorded below.

Parameters

NameType & defaultContract
arrpositional-or-keywordArrangementRequiredSee the operation contract and implementation below.
xpositional-or-keywordTensorRequiredSee the operation contract and implementation below.
signspositional-or-keywordSequence[int]RequiredSee the operation contract and implementation below.
betakeyword-onlyfloatRequiredSee the operation contract and implementation below.
ReturnsTensor
Explicit validation & failure cases
  • When beta <= 0.0, raises ValueError('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" -q
Read 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)