API REFERENCE/omnibias-partition

omnibias.partition.jax.weights

0.1.0a2AlphaPython >=3.10AGPL / commercial

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

import omnibias.partition.jax.weights

Module contract & mathematical conventions

JAX soft partition-of-unity weights (algebraically corresponding twin of the numpy reference).

partition_weights_arrays is the jit / grad / vmap-traceable kernel (raw arrays in, weights out); partition_weights evaluates a ~omnibias.partition._core.params.PartitionParams. Both reproduce omnibias.partition._core.weights.partition_weights bit-for-bit (float64, parity ~1e-9).

Terminology: the split gate sigmoid(beta (W.x - t)) hardens as beta -> inf -- the feasibility / temperature sense of "collapse", distinct from the **founding bias collapse** (the multi-bias delta -> 0 limit to the closed-form derivative sigma^(K-1); see docs/derivatives.md).

Numerical scope: shared coefficients and formulas do not establish universal floating-point identity across backends, devices, compiler transforms, or reduction orders. Compare the chosen path with dtype-appropriate tolerances; even float64 is not a general bitwise-equality guarantee. Algebraic initialization identities likewise need a numerical check for the selected coefficients and reduction order.

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

combine

omnibias.partition.jax.combineSource ↗

combine(weights: Any, region_outputs: Any) -> Any

Blend per-region outputs by partition weights: F = sum_l w_l out_l (jax twin).

Blend region predictions using already-computed partition weights. Scalar expert outputs produce one scalar per sample; vector expert outputs retain their final component axis. The blend is differentiable with respect to both tensors on supported backend paths. This function consumes the supplied weights; it does not independently construct or certify a partition of unity.

Parameters

Name	Type & default	Contract
weightspositional-or-keyword	AnyRequired	Region weights, shape (B, L), where L is the number of regions.
region_outputspositional-or-keyword	AnyRequired	Scalar expert outputs (B, L), returning (B,), or vector expert outputs (B, L, K), returning (B, K). Batch and region axes must agree with weights.

ReturnsAny

Read the implementation

def combine(weights: Any, region_outputs: Any) -> Any:
    r"""Blend per-region outputs by partition weights: ``F = sum_l w_l out_l`` (jax twin)."""
    if jnp.ndim(region_outputs) == 2:  # (n, L) scalar-per-region
        return jnp.einsum("nl,nl->n", weights, region_outputs)
    return jnp.einsum("nl,nlk->nk", weights, region_outputs)

Other public import paths

• omnibias.partition.jax.combine
• omnibias.partition.jax.weights.combine

FUNCTION

partition_weights

omnibias.partition.jax.partition_weightsSource ↗

partition_weights(params: PartitionParams, X: Any, beta: float) -> Any

Soft partition weights (n, 2**depth) for a PartitionParams (float64).

Parameters

Name	Type & default	Contract
paramspositional-or-keyword	PartitionParamsRequired	See the operation contract and implementation below.
Xpositional-or-keyword	AnyRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatRequired	See the operation contract and implementation below.

ReturnsAny

Executed regression · jax matches numpy

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_jax_matches_numpy(beta: float) -> None:
    jax = pytest.importorskip("jax")  # noqa: F841
    from omnibias.partition.jax import partition_weights as jw

    params, X = _sample()
    ref = partition_weights(params, X, beta)
    got = np.asarray(jw(params, X, beta))
    assert np.max(np.abs(got - ref)) < 1e-9

python -m pytest "packages/omnibias-partition/tests/test_parity.py::test_jax_matches_numpy[0.5]" -q

Read the implementation

def partition_weights(params: PartitionParams, X: Any, beta: float) -> Any:
    r"""Soft partition weights ``(n, 2**depth)`` for a :class:`PartitionParams` (float64)."""
    Xv = jnp.asarray(X, dtype=jnp.float64)
    return partition_weights_arrays(
        jnp.asarray(params.W), jnp.asarray(params.t), Xv, float(beta), params.depth
    )

Related APIsomnibias.partition.jax.partition_weights_arrays↗

Other public import paths

• omnibias.partition.jax.partition_weights
• omnibias.partition.jax.weights.partition_weights

FUNCTION

partition_weights_arrays

omnibias.partition.jax.partition_weights_arraysSource ↗

partition_weights_arrays(W: Any, t: Any, X: Any, beta: float, depth: int) -> Any

Region weights (n, 2**depth) from raw arrays -- the jax-traceable kernel.

Use the raw-array path when gate weights or thresholds must learn. It forms sigmoid(beta * (X W^T - t)), chooses the corresponding gate or its complement for each binary region code, and multiplies those factors across splits. There are 2**depth regions. Nonnegative weights sum to one up to floating-point error. Large inverse temperature can saturate gates and suppress gradients; validate the operating range of the model.

Parameters

Name	Type & default	Contract
Wpositional-or-keyword	AnyRequired	Gate coefficient matrix with shape (depth, D). Each row defines one shared binary split.
tpositional-or-keyword	AnyRequired	Threshold vector with shape (depth,).
Xpositional-or-keyword	AnyRequired	Input samples with shape (B, D). Gate tensors must use compatible dtype and device.
betapositional-or-keyword	floatRequired	Finite inverse-temperature scale for the sigmoid gates. Increasing a positive beta sharpens the split; it changes gradient concentration.
depthpositional-or-keyword	intRequired	Number of shared splits. It must match the leading dimension of W and the length of t; the output has shape (B, 2**depth).

ReturnsAny

Executed regression · jax matches numpy

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_jax_matches_numpy(beta: float) -> None:
    jax = pytest.importorskip("jax")  # noqa: F841
    from omnibias.partition.jax import partition_weights as jw

    params, X = _sample()
    ref = partition_weights(params, X, beta)
    got = np.asarray(jw(params, X, beta))
    assert np.max(np.abs(got - ref)) < 1e-9

python -m pytest "packages/omnibias-partition/tests/test_parity.py::test_jax_matches_numpy[0.5]" -q

Read the implementation

def partition_weights_arrays(W: Any, t: Any, X: Any, beta: float, depth: int) -> Any:
    r"""Region weights ``(n, 2**depth)`` from raw arrays -- the ``jax``-traceable kernel."""
    codes = jnp.asarray(region_code_matrix(depth))  # (L, D)
    z = X @ W.T - t[None, :]  # (n, D)
    g = jax.nn.sigmoid(beta * z)  # (n, D)
    gexp = g[:, None, :]  # (n, 1, D)
    bexp = codes[None, :, :]  # (1, L, D)
    factors = bexp * gexp + (1.0 - bexp) * (1.0 - gexp)  # (n, L, D)
    return jnp.prod(factors, axis=-1)  # (n, L)

Related APIsomnibias.partition.region_code_matrix↗

Other public import paths

• omnibias.partition.jax.partition_weights_arrays
• omnibias.partition.jax.weights.partition_weights_arrays

← Package overviewRunnable learning examples →