omnibias.partition.registry

0.1.0a2AlphaPython >=3.10AGPL / commercial

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

import omnibias.partition.registry

Module contract & mathematical conventions

Per-region model registry and the one combine engine every bridge calls.

RegionModels attaches a callable m_l : X -> out to each of the partition's 2**depth regions and blends their outputs by the soft partition-of-unity weights:

This single sum_l w_l * out_l blend is the shared engine behind every downstream bridge (the discontinuity PINN's u = sum_l w_l u_l, the atlas's g = sum_l w_l G_l, the piecewise-symbolic mixture, and the decision layer's relaxed choice). The numpy combine_outputs is the reference; the torch / jax weight twins reproduce the same blend with their own arrays.

Terminology: the weights harden 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).

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 →

CONSTANT

RegionModel

omnibias.partition.registry.RegionModelSource ↗

RegionModel = Callable[[FloatArray], FloatArray]

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

RegionModel = Callable[[FloatArray], FloatArray]

CLASS

RegionModels

omnibias.partition.RegionModelsSource ↗

class RegionModels

Dataclass constructor fields; inherited fields, if any, are declared by the base class. A field(...) expression describes a default factory or field options, not a literal argument default.

A partition plus one callable per region; combine blends them by the POU.

Parameters

Name	Type & default	Contract
`params`dataclass field	`PartitionParams`Required	See the operation contract and implementation below.
`models`dataclass field	`Sequence[RegionModel]`Required	See the operation contract and implementation below.

Executed regression · wrong number of models raises

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_wrong_number_of_models_raises() -> None:
    cfg = PartitionConfig(n_features=2, depth=2, seed=0)  # 4 regions
    params = init_params(cfg)
    with pytest.raises(ValueError, match="region models"):
        RegionModels(params, [lambda X: X[:, :1]])

python -m pytest "packages/omnibias-partition/tests/test_registry.py::test_wrong_number_of_models_raises" -q

Regression example · region models combine is convex blend

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_region_models_combine_is_convex_blend() -> None:
    cfg = PartitionConfig(n_features=2, depth=1, seed=1)
    params = init_params(cfg, rng=1)
    consts = [0.0, 10.0]  # two constant region models
    models = [(lambda X, c=c: np.full((X.shape[0], 1), c)) for c in consts]
    reg = RegionModels(params, models)
    X = np.random.default_rng(2).standard_normal((25, 2))
    out = reg.combine(X, beta=3.0)
    w = partition_weights(params, X, beta=3.0)
    expected = w[:, 0] * consts[0] + w[:, 1] * consts[1]
    assert np.allclose(out, expected)
    assert np.all((out >= -1e-9) & (out <= 10.0 + 1e-9))  # bounded by the region values

python -m pytest packages/omnibias-partition/tests/test_registry.py -k test_region_models_combine_is_convex_blend -q

Fields & methods

paramsfield

params: PartitionParams

Declared PartitionParams field on RegionModels. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

modelsfield

models: Sequence[RegionModel]

Declared Sequence[RegionModel] field on RegionModels. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

__post_init__method

__post_init__(self) -> None

See the class contract above and the exact method implementation below.

Returns: None

Validation & failure cases

When len(self.models) != self.params.n_regions, raises ValueError(f'expected {self.params.n_regions} region models (2**depth), got {len(self.models)}').

Executed regression · wrong number of models raises

Full file, imports & fixtures ↗

def test_wrong_number_of_models_raises() -> None:
    cfg = PartitionConfig(n_features=2, depth=2, seed=0)  # 4 regions
    params = init_params(cfg)
    with pytest.raises(ValueError, match="region models"):
        RegionModels(params, [lambda X: X[:, :1]])

python -m pytest "packages/omnibias-partition/tests/test_registry.py::test_wrong_number_of_models_raises" -q

Method implementation

    def __post_init__(self) -> None:
        if len(self.models) != self.params.n_regions:
            raise ValueError(
                f"expected {self.params.n_regions} region models (2**depth), "
                f"got {len(self.models)}"
            )

View source ↗

region_outputsmethod

region_outputs(self, X: FloatArray) -> FloatArray

Stack each region model's output -> (n, n_regions, k).

Parameters

Name	Type & default	Contract
`X`positional-or-keyword	`FloatArray`Required	See the operation contract and implementation below.

Returns: FloatArray

Executed regression · region models combine is convex blend

Full file, imports & fixtures ↗

def test_region_models_combine_is_convex_blend() -> None:
    cfg = PartitionConfig(n_features=2, depth=1, seed=1)
    params = init_params(cfg, rng=1)
    consts = [0.0, 10.0]  # two constant region models
    models = [(lambda X, c=c: np.full((X.shape[0], 1), c)) for c in consts]
    reg = RegionModels(params, models)
    X = np.random.default_rng(2).standard_normal((25, 2))
    out = reg.combine(X, beta=3.0)
    w = partition_weights(params, X, beta=3.0)
    expected = w[:, 0] * consts[0] + w[:, 1] * consts[1]
    assert np.allclose(out, expected)
    assert np.all((out >= -1e-9) & (out <= 10.0 + 1e-9))  # bounded by the region values

python -m pytest "packages/omnibias-partition/tests/test_registry.py::test_region_models_combine_is_convex_blend" -q

Method implementation

    def region_outputs(self, X: FloatArray) -> FloatArray:
        r"""Stack each region model's output -> ``(n, n_regions, k)``."""
        Xv = np.asarray(X, dtype=np.float64)
        outs = []
        for m in self.models:
            o = np.asarray(m(Xv), dtype=np.float64)
            outs.append(o.reshape(o.shape[0], -1))
        return np.stack(outs, axis=1)  # (n, L, k)

View source ↗

combinemethod

combine(self, X: FloatArray, beta: float | None=None) -> FloatArray

Evaluate F(x) = sum_l w_l(x) m_l(x) at gate sharpness beta.

Parameters

Name	Type & default	Contract
`X`positional-or-keyword	`FloatArray`Required	See the operation contract and implementation below.
`beta`positional-or-keyword	`float \| None`Default: None	See the operation contract and implementation below.

Returns: FloatArray

Executed regression · region models combine is convex blend

Full file, imports & fixtures ↗

def test_region_models_combine_is_convex_blend() -> None:
    cfg = PartitionConfig(n_features=2, depth=1, seed=1)
    params = init_params(cfg, rng=1)
    consts = [0.0, 10.0]  # two constant region models
    models = [(lambda X, c=c: np.full((X.shape[0], 1), c)) for c in consts]
    reg = RegionModels(params, models)
    X = np.random.default_rng(2).standard_normal((25, 2))
    out = reg.combine(X, beta=3.0)
    w = partition_weights(params, X, beta=3.0)
    expected = w[:, 0] * consts[0] + w[:, 1] * consts[1]
    assert np.allclose(out, expected)
    assert np.all((out >= -1e-9) & (out <= 10.0 + 1e-9))  # bounded by the region values

python -m pytest "packages/omnibias-partition/tests/test_registry.py::test_region_models_combine_is_convex_blend" -q

Method implementation

    def combine(self, X: FloatArray, beta: float | None = None) -> FloatArray:
        r"""Evaluate ``F(x) = sum_l w_l(x) m_l(x)`` at gate sharpness ``beta``."""
        b = float(self.params.config.beta_final if beta is None else beta)
        Xv = np.asarray(X, dtype=np.float64)
        w = partition_weights(self.params, Xv, b)  # (n, L)
        outs = self.region_outputs(Xv)  # (n, L, k)
        out = combine_outputs(w, outs)
        return out[:, 0] if (out.ndim == 2 and out.shape[1] == 1) else out

View source ↗

Read the implementation

class RegionModels:
    r"""A partition plus one callable per region; :meth:`combine` blends them by the POU."""

    params: PartitionParams
    models: Sequence[RegionModel]

    def __post_init__(self) -> None:
        if len(self.models) != self.params.n_regions:
            raise ValueError(
                f"expected {self.params.n_regions} region models (2**depth), "
                f"got {len(self.models)}"
            )

    def region_outputs(self, X: FloatArray) -> FloatArray:
        r"""Stack each region model's output -> ``(n, n_regions, k)``."""
        Xv = np.asarray(X, dtype=np.float64)
        outs = []
        for m in self.models:
            o = np.asarray(m(Xv), dtype=np.float64)
            outs.append(o.reshape(o.shape[0], -1))
        return np.stack(outs, axis=1)  # (n, L, k)

    def combine(self, X: FloatArray, beta: float | None = None) -> FloatArray:
        r"""Evaluate ``F(x) = sum_l w_l(x) m_l(x)`` at gate sharpness ``beta``."""
        b = float(self.params.config.beta_final if beta is None else beta)
        Xv = np.asarray(X, dtype=np.float64)
        w = partition_weights(self.params, Xv, b)  # (n, L)
        outs = self.region_outputs(Xv)  # (n, L, k)
        out = combine_outputs(w, outs)
        return out[:, 0] if (out.ndim == 2 and out.shape[1] == 1) else out

Other public import paths

omnibias.partition.RegionModels
omnibias.partition.registry.RegionModels

FUNCTION

combine_outputs

omnibias.partition.combine_outputsSource ↗

combine_outputs(weights: FloatArray, region_outputs: FloatArray) -> FloatArray

Blend per-region outputs by partition weights: F[n, k] = sum_l w[n, l] out[n, l, k].

Parameters

weights: Partition-of-unity weights (n, n_regions) (rows non-negative, summing to one). region_outputs: Per-region outputs (n, n_regions, k) (or (n, n_regions) for scalar outputs).

Parameters

Name	Type & default	Contract
`weights`positional-or-keyword	`FloatArray`Required	Partition-of-unity weights `(n, n_regions)` (rows non-negative, summing to one).
`region_outputs`positional-or-keyword	`FloatArray`Required	Per-region outputs `(n, n_regions, k)` (or `(n, n_regions)` for scalar outputs).

ReturnsFloatArray

Executed regression · combine outputs matches manual blend

Full file, imports & fixtures ↗

def test_combine_outputs_matches_manual_blend() -> None:
    rng = np.random.default_rng(0)
    w = rng.uniform(size=(10, 4))
    w = w / w.sum(axis=1, keepdims=True)  # a partition of unity
    outs = rng.standard_normal((10, 4, 3))
    got = combine_outputs(w, outs)
    ref = np.einsum("nl,nlk->nk", w, outs)
    assert np.allclose(got, ref)

python -m pytest "packages/omnibias-partition/tests/test_registry.py::test_combine_outputs_matches_manual_blend" -q

Read the implementation

def combine_outputs(weights: FloatArray, region_outputs: FloatArray) -> FloatArray:
    r"""Blend per-region outputs by partition weights: ``F[n, k] = sum_l w[n, l] out[n, l, k]``.

    Parameters
    ----------
    weights:
        Partition-of-unity weights ``(n, n_regions)`` (rows non-negative, summing to one).
    region_outputs:
        Per-region outputs ``(n, n_regions, k)`` (or ``(n, n_regions)`` for scalar outputs).
    """
    w = np.asarray(weights, dtype=np.float64)
    outs = np.asarray(region_outputs, dtype=np.float64)
    if outs.ndim == 2:  # scalar per region -> (n, n_regions)
        return np.einsum("nl,nl->n", w, outs)
    return np.einsum("nl,nlk->nk", w, outs)

Other public import paths

omnibias.partition.combine_outputs
omnibias.partition.registry.combine_outputs

← Package overview Runnable learning examples →

omnibias.partition.registry

RegionModel​

RegionModels​

Parameters

Fields & methods

Parameters

Parameters

combine_outputs​

Parameters

RegionModel

RegionModels

combine_outputs