API REFERENCE/omnibias-partition

omnibias.partition.keras.weights

0.1.0a2AlphaPython >=3.10AGPL / commercial

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

import omnibias.partition.keras.weights

Module contract & mathematical conventions

Keras 3 soft partition-of-unity weights (keras.ops twin of numpy / torch / jax).

Parity with omnibias.partition._core.weights.partition_weights is ~1e-9 in float64. Lives here (not in omnibias-keras) so the keras core package never depends on partition / tab.

Terminology: the split gate sigmoid(beta (W.x - t)) hardens as beta -> inf (the feasibility / temperature sense of collapse), distinct from the founding delta -> 0 bias collapse.

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.keras.combineSource ↗

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

Blend per-region outputs: F = sum_l w_l out_l.

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: ``F = sum_l w_l out_l``."""
    if ops.ndim(region_outputs) == 2:
        return ops.einsum("nl,nl->n", weights, region_outputs)
    return ops.einsum("nl,nlk->nk", weights, region_outputs)

Other public import paths

• omnibias.partition.keras.combine
• omnibias.partition.keras.weights.combine

FUNCTION

partition_weights

omnibias.partition.keras.partition_weightsSource ↗

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

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

Parameters

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

ReturnsAny

Read the implementation

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

Related APIsomnibias.partition.keras.partition_weights_arrays↗

Other public import paths

• omnibias.partition.keras.partition_weights
• omnibias.partition.keras.weights.partition_weights

FUNCTION

partition_weights_arrays

omnibias.partition.keras.partition_weights_arraysSource ↗

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

Region weights (n, 2**depth) from raw tensors -- the keras.ops 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 · keras partition weights 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_keras_partition_weights_parity() -> None:
    from omnibias.partition import PartitionConfig
    from omnibias.partition._core.params import PartitionParams
    from omnibias.partition._core.weights import partition_weights as partition_np
    from omnibias.partition.keras.weights import partition_weights_arrays

    rng = np.random.default_rng(3)
    X = rng.standard_normal((16, 4))
    W = rng.standard_normal((2, 4)) * 0.4
    t = rng.standard_normal((2,)) * 0.2
    beta = 5.0
    cfg = PartitionConfig(n_features=4, depth=2, split_kind="oblique", beta_final=beta)
    params = PartitionParams(cfg, W, t)
    np_w = partition_np(params, X, beta)
    ops = keras.ops
    keras_w = _np(
        partition_weights_arrays(
            ops.convert_to_tensor(W.astype(np.float64), dtype="float64"),
            ops.convert_to_tensor(t.astype(np.float64), dtype="float64"),
            ops.convert_to_tensor(X.astype(np.float64), dtype="float64"),
            beta,
            2,
        )
    )
    assert keras_w.dtype == np.float64
    assert np.max(np.abs(np_w - keras_w)) < _atol()

python -m pytest "packages/omnibias-partition/tests/test_keras_weights.py::test_keras_partition_weights_parity" -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 tensors -- the ``keras.ops`` kernel."""
    codes = ops.convert_to_tensor(region_code_matrix(int(depth)), dtype=X.dtype)
    z = ops.matmul(X, ops.transpose(W)) - ops.expand_dims(t, 0)
    g = ops.sigmoid(ops.cast(beta, X.dtype) * z)
    gexp = ops.expand_dims(g, 1)
    bexp = ops.expand_dims(codes, 0)
    factors = bexp * gexp + (1.0 - bexp) * (1.0 - gexp)
    return prod_last_axis(factors, int(depth))

Related APIsomnibias.partition.region_code_matrix↗omnibias.partition.keras.prod_last_axis↗

Other public import paths

• omnibias.partition.keras.partition_weights_arrays
• omnibias.partition.keras.weights.partition_weights_arrays

FUNCTION

prod_last_axis

omnibias.partition.keras.prod_last_axisSource ↗

prod_last_axis(x: Any, size: int) -> Any

Product over the last axis, preserving operand dtype.

keras.ops.prod on the torch backend returns float32 even when x is float64, which breaks float64 parity. Sequential multiply keeps dtype.

Parameters

Name	Type & default	Contract
xpositional-or-keyword	AnyRequired	See the operation contract and implementation below.
sizepositional-or-keyword	intRequired	See the operation contract and implementation below.

ReturnsAny

Explicit validation & failure cases

• When n < 1, raises ValueError('size must be >= 1').

These are explicit source guards. Backend operations may also raise for invalid shapes, dtypes, or unsupported devices.

Executed regression · keras partition weights 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_keras_partition_weights_parity() -> None:
    from omnibias.partition import PartitionConfig
    from omnibias.partition._core.params import PartitionParams
    from omnibias.partition._core.weights import partition_weights as partition_np
    from omnibias.partition.keras.weights import partition_weights_arrays

    rng = np.random.default_rng(3)
    X = rng.standard_normal((16, 4))
    W = rng.standard_normal((2, 4)) * 0.4
    t = rng.standard_normal((2,)) * 0.2
    beta = 5.0
    cfg = PartitionConfig(n_features=4, depth=2, split_kind="oblique", beta_final=beta)
    params = PartitionParams(cfg, W, t)
    np_w = partition_np(params, X, beta)
    ops = keras.ops
    keras_w = _np(
        partition_weights_arrays(
            ops.convert_to_tensor(W.astype(np.float64), dtype="float64"),
            ops.convert_to_tensor(t.astype(np.float64), dtype="float64"),
            ops.convert_to_tensor(X.astype(np.float64), dtype="float64"),
            beta,
            2,
        )
    )
    assert keras_w.dtype == np.float64
    assert np.max(np.abs(np_w - keras_w)) < _atol()

python -m pytest "packages/omnibias-partition/tests/test_keras_weights.py::test_keras_partition_weights_parity" -q

Read the implementation

def prod_last_axis(x: Any, size: int) -> Any:
    """Product over the last axis, preserving operand dtype.

    ``keras.ops.prod`` on the torch backend returns float32 even when ``x`` is
    float64, which breaks float64 parity. Sequential multiply keeps dtype.
    """
    n = int(size)
    if n < 1:
        raise ValueError("size must be >= 1")
    acc = x[..., 0]
    for i in range(1, n):
        acc = acc * x[..., i]
    return acc

Other public import paths

• omnibias.partition.keras.prod_last_axis
• omnibias.partition.keras.weights.prod_last_axis

← Package overviewRunnable learning examples →