Skip to main content
API REFERENCE/omnibias-partition

omnibias.partition.arrangement.jax

0.1.0a2AlphaPython >=3.10AGPL / commercial

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

import omnibias.partition.arrangement.jax
Module contract & mathematical conventions

Soft arrangement membership (jax; 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.

Module source ↗Source revision 9a2d73df0573Running reference examples →
FUNCTION

margin​

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

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-keywordArrayRequiredSee the operation contract and implementation below.
ReturnsArray
Read the implementation
def margin(arr: Arrangement, x: Array) -> Array:
    w = jnp.asarray(arr.normals, dtype=x.dtype)
    t = jnp.asarray(arr.offsets, dtype=x.dtype)
    xv = x if x.ndim == 2 else x[None, :]
    z = xv @ w.T - t
    return jnp.min(jnp.abs(z), axis=-1)
FUNCTION

soft_membership​

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

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-keywordArrayRequiredSee 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.
ReturnsArray
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.

Full file, imports & fixtures ↗
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
Regression example · g3 tree agrees with partition 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.

Full file, imports & fixtures ↗
def test_g3_tree_agrees_with_partition_weights() -> None:
    cfg = PartitionConfig(n_features=3, depth=3)
    params = init_params(cfg, rng=0)
    arr = tree_arrangement(params.W, params.t)
    x = np.array([[0.2, -0.1, 0.4], [1.0, 0.0, -0.5], [-0.3, 0.7, 0.1]], dtype=np.float64)
    beta = 3.5
    pw = partition_weights(params, x, beta)
    codes = region_code_matrix(3)
    eps = np.finfo(np.float64).eps
    for leaf in range(8):
        signs = tuple(1 if codes[leaf, j] > 0.5 else -1 for j in range(3))
        sm = soft_membership(arr, x, signs, beta=beta)
        for a, b in zip(sm, pw[:, leaf], strict=True):
            scale = max(abs(float(a)), abs(float(b)), 1.0)
            ulp = abs(float(a) - float(b)) / (eps * scale)
            assert ulp <= 4.0
python -m pytest packages/omnibias-partition/tests/test_arrangement.py -k test_g3_tree_agrees_with_partition_weights -q
Read the implementation
def soft_membership(
    arr: Arrangement,
    x: Array,
    signs: Sequence[int],
    *,
    beta: float,
) -> Array:
    if beta <= 0.0:
        raise ValueError("beta must be > 0 (temperature collapse axis)")
    w = jnp.asarray(arr.normals, dtype=x.dtype)
    t = jnp.asarray(arr.offsets, dtype=x.dtype)
    xv = x if x.ndim == 2 else x[None, :]
    z = xv @ w.T - t
    s = jnp.asarray(list(signs), dtype=x.dtype)
    return sigmoid(beta * s * z).prod(axis=-1)
← Package overviewRunnable learning examples →