API REFERENCE/omnibias-partition

omnibias.partition

0.1.0a2AlphaPython >=3.10AGPL / commercial

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

import omnibias.partition

Choose partition to combine local experts using nonnegative region weights that sum to one under the supported gate construction. The raw differentiable array APIs are useful when gate matrices and thresholds are trainable tensors; configuration wrappers may have different conversion and gradient behavior.

A partition tree's depth determines its number of leaves and the required expert axis. Keep batch, gate, feature, and leaf dimensions explicit. The soft regional-model example checks the partition-of-unity identity and verifies that gradients reach gate weights, thresholds, and experts.

Increasing inverse temperature sharpens region membership, but can saturate gradients. A finite-temperature membership-gap certificate has a specific margin and input-domain scope; it is not a global guarantee about an arbitrary learned predictor. This package uses AGPL-or-commercial licensing.

Module contract & mathematical conventions

omnibias-partition: a light, certified soft partition-of-unity primitive.

A hard partition of R^d into regions is a set of indicator functions 1[region l]. omnibias makes it a **soft partition** built from oblique split gates g(x) = sigmoid(beta * (w.x - t)): depth gates route an input into 2**depth regions with weights w_l(x) that are non-negative, sum to one (a genuine partition of unity), and **harden** to a crisp {0, 1} partition as beta -> inf.

The keystone shared by four downstream bridges (a discontinuity-capturing PINN, a region-wise Riemannian atlas, per-region symbolic discovery, and a certified decision layer). It ships:

1. partition_weights -- the numpy reference plus algebraically corresponding omnibias.partition.torch / omnibias.partition.jax twins (parity ~1e-9); 2. hard_assignment / hardened_rules -- the crisp region index and the exported human-readable if w.x > t boundaries; 3. certify_partition_gap -- a **sound** soft->hard membership-gap certificate (outward-rounded ~omnibias.core.verified.Interval + closed-form log(n_regions)/beta), a well-posed **yes-if** object; 4. RegionModels -- a per-region model registry whose single combine(X, beta) = sum_l w_l * out_l engine every bridge calls.

Terminology: the gate's beta -> inf hardening is **temperature collapse** -- the feasibility sense (a soft indicator becoming a 0/1 step), distinct from the **founding bias collapse** (the multi-bias delta -> 0 limit of an OMBU to the closed-form derivative sigma^(K-1); see docs/derivatives.md). The bridges differentiate products of sigmoids by autodiff -- the closed-form derivative tower does not auto-extend to products.

The torch / jax weight twins need the torch / jax extra (each degrades gracefully; the numpy core + certificate need neither).

important:

**Bit-parity with the PyTorch twin requires 64-bit JAX** -- jax.config.update("jax_enable_x64", True) before the first JAX array is created (or JAX_ENABLE_X64=1). JAX otherwise truncates to float32 while PyTorch uses float64, so the twins stay internally consistent but agree only to float32 tolerance. Where a value feeds a threshold, a rounding step or an argmax, that is enough to change the decision rather than just the last digits. See omnibias.jax.precision.

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 →

arrangementModule ↗certifyModule ↗jaxModule ↗kerasModule ↗registryModule ↗torchModule ↗

Imports & aliases

These names share a definition with another public module. Follow a name for its full contract, parameters, and examples.

Import from this module	Canonical reference
PartitionGapCertificate	omnibias.partition.PartitionGapCertificate
RegionModels	omnibias.partition.RegionModels
certify_partition_gap	omnibias.partition.certify_partition_gap
combine_outputs	omnibias.partition.combine_outputs

CLASS

PartitionConfig

omnibias.partition.PartitionConfigSource ↗

class PartitionConfig

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.

Shape + split-geometry descriptor for a soft partition of unity.

A single soft tree of depth oblique gates routes R^{n_features} into 2**depth regions; the region weights are a partition of unity (non-negative, sum to one) that hardens to a crisp partition as beta -> inf.

Attributes

n_features: Input dimension d. depth: Number of split gates. The partition has 2**depth regions. split_kind: "oblique" (dense hyperplanes), "axis" (single-feature thresholds) or "sparse" (oblique storage, L1-sparsified by a trainer). beta_init, beta_final: Gate sharpness (inverse temperature) at the start / end of the anneal. The beta -> inf limit is the feasibility-sense collapse to a hard partition. anneal_steps: Steps over which beta ramps beta_init -> beta_final (geometric). seed: RNG seed for parameter initialisation.

Parameters

Name	Type & default	Contract
n_featuresdataclass field	intRequired	Input dimension d.
depthdataclass field	intDefault: 1	Number of split gates. The partition has 2**depth regions.
split_kinddataclass field	strDefault: 'oblique'	"oblique" (dense hyperplanes), "axis" (single-feature thresholds) or "sparse" (oblique storage, L1-sparsified by a trainer).
beta_initdataclass field	floatDefault: 1.0	Gate sharpness (inverse temperature) at the start / end of the anneal. The beta -> inf limit is the feasibility-sense collapse to a hard partition.
beta_finaldataclass field	floatDefault: 32.0	See the operation contract and implementation below.
anneal_stepsdataclass field	intDefault: 50	Steps over which beta ramps beta_init -> beta_final (geometric).
seeddataclass field	intDefault: 0	RNG seed for parameter initialisation.

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 · axis init is single feature

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_axis_init_is_single_feature() -> None:
    cfg = PartitionConfig(n_features=5, depth=3, split_kind="axis", seed=7)
    params = init_params(cfg)
    for j in range(cfg.depth):
        assert np.count_nonzero(params.W[j]) == 1  # each gate reads one feature

python -m pytest packages/omnibias-partition/tests/test_weights.py -k test_axis_init_is_single_feature -q

Fields & methods

n_featuresfield

n_features: int

Input dimension d.

View source ↗

depthfield

depth: int = 1

Number of split gates. The partition has 2**depth regions.

Regression example · axis init is single feature

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_axis_init_is_single_feature() -> None:
    cfg = PartitionConfig(n_features=5, depth=3, split_kind="axis", seed=7)
    params = init_params(cfg)
    for j in range(cfg.depth):
        assert np.count_nonzero(params.W[j]) == 1  # each gate reads one feature

python -m pytest packages/omnibias-partition/tests/test_weights.py -k test_axis_init_is_single_feature -q

View source ↗

split_kindfield

split_kind: str = 'oblique'

"oblique" (dense hyperplanes), "axis" (single-feature thresholds) or "sparse" (oblique storage, L1-sparsified by a trainer).

View source ↗

beta_initfield

beta_init: float = 1.0

Gate sharpness (inverse temperature) at the start / end of the anneal. The beta -> inf limit is the feasibility-sense collapse to a hard partition.

View source ↗

beta_finalfield

beta_final: float = 32.0

Gate sharpness (inverse temperature) at the start / end of the anneal. The beta -> inf limit is the feasibility-sense collapse to a hard partition.

View source ↗

anneal_stepsfield

anneal_steps: int = 50

Steps over which beta ramps beta_init -> beta_final (geometric).

View source ↗

seedfield

seed: int = 0

RNG seed for parameter initialisation.

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 self.n_features < 1, raises ValueError(f'n_features must be >= 1, got {self.n_features}').
• When self.depth < 1, raises ValueError(f'depth must be >= 1, got {self.depth}').
• When self.split_kind not in _SPLIT_KINDS, raises ValueError(f'split_kind must be one of {_SPLIT_KINDS}, got {self.split_kind!r}').
• When not (self.beta_init > 0.0 and self.beta_final > 0.0), raises ValueError('beta_init and beta_final must be positive').
• When self.anneal_steps < 1, raises ValueError('anneal_steps must be >= 1').

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

Method implementation

    def __post_init__(self) -> None:
        if self.n_features < 1:
            raise ValueError(f"n_features must be >= 1, got {self.n_features}")
        if self.depth < 1:
            raise ValueError(f"depth must be >= 1, got {self.depth}")
        if self.split_kind not in _SPLIT_KINDS:
            raise ValueError(f"split_kind must be one of {_SPLIT_KINDS}, got {self.split_kind!r}")
        if not (self.beta_init > 0.0 and self.beta_final > 0.0):
            raise ValueError("beta_init and beta_final must be positive")
        if self.anneal_steps < 1:
            raise ValueError("anneal_steps must be >= 1")

View source ↗

n_regionsproperty

n_regions(self) -> int

Number of regions, 2**depth.

Returns: int

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

Method implementation

    def n_regions(self) -> int:
        r"""Number of regions, ``2**depth``."""
        return 1 << self.depth

View source ↗

is_axisproperty

is_axis(self) -> bool

True for the single-feature (axis-aligned) split mode.

Returns: bool

Executed regression · hardened rules are readable

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_hardened_rules_are_readable() -> None:
    cfg = PartitionConfig(n_features=3, depth=2, split_kind="axis", seed=1)
    params = init_params(cfg)
    rules = hardened_rules(params)
    assert len(rules) == cfg.depth
    assert all("x[" in r for r in rules)  # single-feature threshold form
    clause = region_rule(params, region=cfg.n_regions - 1)
    assert "AND" in clause

python -m pytest "packages/omnibias-partition/tests/test_weights.py::test_hardened_rules_are_readable" -q

Method implementation

    def is_axis(self) -> bool:
        r"""``True`` for the single-feature (axis-aligned) split mode."""
        return self.split_kind == "axis"

View source ↗

beta_atmethod

beta_at(self, step: int) -> float

Geometric beta ramp: beta_init at step 0 -> beta_final at the end.

Parameters

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

Returns: float

Method implementation

    def beta_at(self, step: int) -> float:
        r"""Geometric ``beta`` ramp: ``beta_init`` at step 0 -> ``beta_final`` at the end."""
        if self.anneal_steps <= 1:
            return float(self.beta_final)
        frac = min(max(step, 0), self.anneal_steps - 1) / (self.anneal_steps - 1)
        return float(self.beta_init * (self.beta_final / self.beta_init) ** frac)

View source ↗

Read the implementation

class PartitionConfig:
    r"""Shape + split-geometry descriptor for a soft partition of unity.

    A single soft tree of ``depth`` oblique gates routes ``R^{n_features}`` into
    ``2**depth`` regions; the region weights are a partition of unity (non-negative, sum to
    one) that hardens to a crisp partition as ``beta -> inf``.

    Attributes
    ----------
    n_features:
        Input dimension ``d``.
    depth:
        Number of split gates. The partition has ``2**depth`` regions.
    split_kind:
        ``"oblique"`` (dense hyperplanes), ``"axis"`` (single-feature thresholds) or
        ``"sparse"`` (oblique storage, L1-sparsified by a trainer).
    beta_init, beta_final:
        Gate sharpness (inverse temperature) at the start / end of the anneal. The
        ``beta -> inf`` limit is the feasibility-sense collapse to a hard partition.
    anneal_steps:
        Steps over which ``beta`` ramps ``beta_init -> beta_final`` (geometric).
    seed:
        RNG seed for parameter initialisation.
    """

    n_features: int
    depth: int = 1
    split_kind: str = "oblique"
    beta_init: float = 1.0
    beta_final: float = 32.0
    anneal_steps: int = 50
    seed: int = 0

    def __post_init__(self) -> None:
        if self.n_features < 1:
            raise ValueError(f"n_features must be >= 1, got {self.n_features}")
        if self.depth < 1:
            raise ValueError(f"depth must be >= 1, got {self.depth}")
        if self.split_kind not in _SPLIT_KINDS:
            raise ValueError(f"split_kind must be one of {_SPLIT_KINDS}, got {self.split_kind!r}")
        if not (self.beta_init > 0.0 and self.beta_final > 0.0):
            raise ValueError("beta_init and beta_final must be positive")
        if self.anneal_steps < 1:
            raise ValueError("anneal_steps must be >= 1")

    @property
    def n_regions(self) -> int:
        r"""Number of regions, ``2**depth``."""
        return 1 << self.depth

    @property
    def is_axis(self) -> bool:
        r"""``True`` for the single-feature (axis-aligned) split mode."""
        return self.split_kind == "axis"

    def beta_at(self, step: int) -> float:
        r"""Geometric ``beta`` ramp: ``beta_init`` at step 0 -> ``beta_final`` at the end."""
        if self.anneal_steps <= 1:
            return float(self.beta_final)
        frac = min(max(step, 0), self.anneal_steps - 1) / (self.anneal_steps - 1)
        return float(self.beta_init * (self.beta_final / self.beta_init) ** frac)

CLASS

PartitionParams

omnibias.partition.PartitionParamsSource ↗

class PartitionParams

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.

The numpy arrays defining a soft partition (see the module docstring).

Parameters

Name	Type & default	Contract
configdataclass field	PartitionConfigRequired	See the operation contract and implementation below.
Wdataclass field	FloatArrayRequired	See the operation contract and implementation below.
tdataclass field	FloatArrayRequired	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 · keras partition weights parity

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 -k test_keras_partition_weights_parity -q

Fields & methods

configfield

config: PartitionConfig

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

View source ↗

Wfield

W: FloatArray

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

View source ↗

tfield

t: FloatArray

Declared FloatArray field on PartitionParams. 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

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

Method implementation

    def __post_init__(self) -> None:
        cfg = self.config
        self.W = np.asarray(self.W, dtype=np.float64).reshape(cfg.depth, cfg.n_features)
        self.t = np.asarray(self.t, dtype=np.float64).reshape(cfg.depth)

View source ↗

n_featuresproperty

n_features(self) -> int

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

Returns: int

Executed regression · interval weight bounds prove containment far from boundary

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_interval_weight_bounds_prove_containment_far_from_boundary() -> None:
    # a box entirely on one side of a single axis split -> that region's weight ~ 1
    cfg = PartitionConfig(n_features=1, depth=1, split_kind="axis", seed=0)
    params = init_params(cfg)
    params.W[:] = 1.0
    params.t[:] = 0.0
    box = np.array([[5.0], [10.0]])  # x in [5, 10], well past the x > 0 split
    bounds = interval_weight_bounds(params, box, beta=50.0)
    # region 1 ("gate fired") is certified to carry essentially all the mass
    assert bounds[1].lo > 0.99

python -m pytest "packages/omnibias-partition/tests/test_certify.py::test_interval_weight_bounds_prove_containment_far_from_boundary" -q

Method implementation

    def n_features(self) -> int:
        return int(self.config.n_features)

View source ↗

depthproperty

depth(self) -> int

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

Returns: int

Executed regression · torch 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_torch_matches_numpy(beta: float) -> None:
    from omnibias.partition.torch import partition_weights as tw

    params, X = _sample()
    ref = partition_weights(params, X, beta)
    got = tw(params, X, beta).detach().cpu().numpy()
    assert np.max(np.abs(got - ref)) < 1e-9

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

Method implementation

    def depth(self) -> int:
        return int(self.config.depth)

View source ↗

n_regionsproperty

n_regions(self) -> int

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

Returns: int

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

Method implementation

    def n_regions(self) -> int:
        return int(self.config.n_regions)

View source ↗

copymethod

copy(self) -> PartitionParams

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

Returns: PartitionParams

Method implementation

    def copy(self) -> PartitionParams:
        return PartitionParams(self.config, self.W.copy(), self.t.copy())

View source ↗

Read the implementation

class PartitionParams:
    r"""The numpy arrays defining a soft partition (see the module docstring)."""

    config: PartitionConfig
    W: FloatArray
    t: FloatArray

    def __post_init__(self) -> None:
        cfg = self.config
        self.W = np.asarray(self.W, dtype=np.float64).reshape(cfg.depth, cfg.n_features)
        self.t = np.asarray(self.t, dtype=np.float64).reshape(cfg.depth)

    @property
    def n_features(self) -> int:
        return int(self.config.n_features)

    @property
    def depth(self) -> int:
        return int(self.config.depth)

    @property
    def n_regions(self) -> int:
        return int(self.config.n_regions)

    def copy(self) -> PartitionParams:
        return PartitionParams(self.config, self.W.copy(), self.t.copy())

FUNCTION

init_params

omnibias.partition.init_paramsSource ↗

init_params(config: PartitionConfig, rng: np.random.Generator | int | None=None, *, weight_scale: float | None=None) -> PartitionParams

Small random initialisation of the depth split gates.

oblique / sparse draw dense directions N(0, weight_scale**2) (default 1/sqrt(n_features)); axis assigns each gate a single feature with a unit direction (an interpretable x[f] > t split). Thresholds start at 0.

Parameters

Name	Type & default	Contract
configpositional-or-keyword	PartitionConfigRequired	See the operation contract and implementation below.
rngpositional-or-keyword	np.random.Generator | int | NoneDefault: None	See the operation contract and implementation below.
weight_scalekeyword-only	float | NoneDefault: None	See the operation contract and implementation below.

ReturnsPartitionParams

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 · axis init is single feature

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_axis_init_is_single_feature() -> None:
    cfg = PartitionConfig(n_features=5, depth=3, split_kind="axis", seed=7)
    params = init_params(cfg)
    for j in range(cfg.depth):
        assert np.count_nonzero(params.W[j]) == 1  # each gate reads one feature

python -m pytest packages/omnibias-partition/tests/test_weights.py -k test_axis_init_is_single_feature -q

Read the implementation

def init_params(
    config: PartitionConfig,
    rng: np.random.Generator | int | None = None,
    *,
    weight_scale: float | None = None,
) -> PartitionParams:
    r"""Small random initialisation of the ``depth`` split gates.

    ``oblique`` / ``sparse`` draw dense directions ``N(0, weight_scale**2)`` (default
    ``1/sqrt(n_features)``); ``axis`` assigns each gate a single feature with a unit
    direction (an interpretable ``x[f] > t`` split). Thresholds start at ``0``.
    """
    gen = _as_rng(rng if rng is not None else config.seed)
    D, d = config.depth, config.n_features
    ws = weight_scale if weight_scale is not None else 1.0 / np.sqrt(d)
    if config.split_kind == "axis":
        W = np.zeros((D, d), dtype=np.float64)
        for j in range(D):
            W[j, int(gen.integers(0, d))] = 1.0
    else:
        W = gen.standard_normal((D, d)) * ws
    t = np.zeros(D, dtype=np.float64)
    return PartitionParams(config, W, t)

Related APIsomnibias.partition.PartitionParams↗

FUNCTION

region_code_matrix

omnibias.partition.region_code_matrixSource ↗

region_code_matrix(depth: int) -> FloatArray

(2**depth, depth) matrix of region bit-codes in {0.0, 1.0}.

Row l is the base-2 expansion of the region index l (bit j = gate j); 1 means "gate j fired" (W.x > t). Shared by every backend so the region ordering is identical.

Parameters

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

ReturnsFloatArray

Executed regression · torch 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_torch_matches_numpy(beta: float) -> None:
    from omnibias.partition.torch import partition_weights as tw

    params, X = _sample()
    ref = partition_weights(params, X, beta)
    got = tw(params, X, beta).detach().cpu().numpy()
    assert np.max(np.abs(got - ref)) < 1e-9

python -m pytest "packages/omnibias-partition/tests/test_parity.py::test_torch_matches_numpy[0.5]" -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 region_code_matrix(depth: int) -> FloatArray:
    r"""``(2**depth, depth)`` matrix of region bit-codes in ``{0.0, 1.0}``.

    Row ``l`` is the base-2 expansion of the region index ``l`` (bit ``j`` = gate ``j``);
    ``1`` means "gate ``j`` fired" (``W.x > t``). Shared by every backend so the region
    ordering is identical.
    """
    L = 1 << depth
    codes = np.zeros((L, depth), dtype=np.float64)
    for region in range(L):
        for j in range(depth):
            codes[region, j] = float((region >> j) & 1)
    return codes

FUNCTION

gate_activations

omnibias.partition.gate_activationsSource ↗

gate_activations(params: PartitionParams, X: FloatArray, beta: float) -> FloatArray

Soft gate activations G of shape (n, depth) in (0, 1).

Parameters

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

ReturnsFloatArray

Executed regression · gate activations shape and range

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_gate_activations_shape_and_range() -> None:
    cfg = PartitionConfig(n_features=2, depth=2, seed=0)
    params = init_params(cfg)
    g = gate_activations(params, np.zeros((5, 2)), beta=1.0)
    assert g.shape == (5, 2)
    assert np.all((g > 0.0) & (g < 1.0))

python -m pytest "packages/omnibias-partition/tests/test_weights.py::test_gate_activations_shape_and_range" -q

Read the implementation

def gate_activations(params: PartitionParams, X: FloatArray, beta: float) -> FloatArray:
    r"""Soft gate activations ``G`` of shape ``(n, depth)`` in ``(0, 1)``."""
    Xv = np.asarray(X, dtype=np.float64)
    z = Xv @ params.W.T - params.t[None, :]  # (n, depth)
    return sigmoid_np(beta * z)

FUNCTION

hard_assignment

omnibias.partition.hard_assignmentSource ↗

hard_assignment(params: PartitionParams, X: FloatArray) -> FloatArray

Crisp region index (n,) from the beta -> inf hard gates 1[W.x > t].

Parameters

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

ReturnsFloatArray

Executed regression · hard assignment selects max weight region when sharp

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_hard_assignment_selects_max_weight_region_when_sharp() -> None:
    cfg = PartitionConfig(n_features=3, depth=2, seed=5)
    params = init_params(cfg, rng=2)
    X = np.random.default_rng(6).standard_normal((40, 3))
    z = np.abs(X @ params.W.T - params.t[None, :])
    interior = np.all(z > 0.05, axis=1)
    soft = partition_weights(params, X, beta=200.0)
    idx = hard_assignment(params, X)
    assert np.array_equal(np.argmax(soft[interior], axis=1), idx[interior])

python -m pytest "packages/omnibias-partition/tests/test_weights.py::test_hard_assignment_selects_max_weight_region_when_sharp" -q

Read the implementation

def hard_assignment(params: PartitionParams, X: FloatArray) -> FloatArray:
    r"""Crisp region index ``(n,)`` from the ``beta -> inf`` hard gates ``1[W.x > t]``."""
    Xv = np.asarray(X, dtype=np.float64)
    z = Xv @ params.W.T - params.t[None, :]  # (n, depth)
    bits = (z > 0.0).astype(np.int64)  # (n, depth)
    weights = 1 << np.arange(params.depth, dtype=np.int64)  # bit j has value 2**j
    return (bits * weights[None, :]).sum(axis=1).astype(np.int64)

FUNCTION

hard_weights

omnibias.partition.hard_weightsSource ↗

hard_weights(params: PartitionParams, X: FloatArray) -> FloatArray

One-hot crisp weights (n, 2**depth) -- the beta -> inf limit of the POU.

Parameters

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

ReturnsFloatArray

Executed regression · certificate is sound and hardens

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_certificate_is_sound_and_hardens() -> None:
    cfg = PartitionConfig(n_features=3, depth=2, seed=4)
    params = init_params(cfg, rng=2)
    X = np.random.default_rng(5).standard_normal((150, 3))
    lo = certify_partition_gap(params, X, beta=2.0)
    hi = certify_partition_gap(params, X, beta=64.0)
    assert lo.is_sound and hi.is_sound
    assert hi.max_gap <= lo.max_gap + 1e-9  # sharper beta -> smaller certified gap
    assert hi.gibbs_scale < lo.gibbs_scale  # log(n_regions)/beta shrinks with beta

python -m pytest "packages/omnibias-partition/tests/test_certify.py::test_certificate_is_sound_and_hardens" -q

Regression example · hardening matches hard 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_hardening_matches_hard_weights() -> None:
    cfg = PartitionConfig(n_features=3, depth=2, seed=3)
    params = init_params(cfg, rng=1)
    X = np.random.default_rng(4).standard_normal((50, 3))
    soft = partition_weights(params, X, beta=400.0)
    hard = hard_weights(params, X)
    # away from the split surface the soft POU collapses onto the hard one-hot
    z = np.abs(X @ params.W.T - params.t[None, :])
    interior = np.all(z > 0.05, axis=1)
    assert np.max(np.abs(soft[interior] - hard[interior])) < 1e-3

python -m pytest packages/omnibias-partition/tests/test_weights.py -k test_hardening_matches_hard_weights -q

Read the implementation

def hard_weights(params: PartitionParams, X: FloatArray) -> FloatArray:
    r"""One-hot crisp weights ``(n, 2**depth)`` -- the ``beta -> inf`` limit of the POU."""
    Xv = np.asarray(X, dtype=np.float64)
    idx = hard_assignment(params, Xv)
    out = np.zeros((Xv.shape[0], params.n_regions), dtype=np.float64)
    out[np.arange(Xv.shape[0]), idx] = 1.0
    return out

Related APIsomnibias.partition.hard_assignment↗

FUNCTION

hardened_rules

omnibias.partition.hardened_rulesSource ↗

hardened_rules(params: PartitionParams) -> list[str]

The exported if ... split boundaries, one per gate (the beta -> inf rules).

Parameters

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

Returnslist[str]

Executed regression · hardened rules are readable

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_hardened_rules_are_readable() -> None:
    cfg = PartitionConfig(n_features=3, depth=2, split_kind="axis", seed=1)
    params = init_params(cfg)
    rules = hardened_rules(params)
    assert len(rules) == cfg.depth
    assert all("x[" in r for r in rules)  # single-feature threshold form
    clause = region_rule(params, region=cfg.n_regions - 1)
    assert "AND" in clause

python -m pytest "packages/omnibias-partition/tests/test_weights.py::test_hardened_rules_are_readable" -q

Read the implementation

def hardened_rules(params: PartitionParams) -> list[str]:
    r"""The exported ``if ...`` split boundaries, one per gate (the ``beta -> inf`` rules)."""
    return [_gate_rule(params, j) for j in range(params.depth)]

FUNCTION

partition_weights

omnibias.partition.partition_weightsSource ↗

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

Soft partition-of-unity weights P of shape (n, 2**depth).

Each row is a probability vector over the regions (non-negative, sums to 1); P[i, l] is the soft-AND membership of sample i in region l.

Parameters

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

ReturnsFloatArray

Executed regression · torch 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_torch_matches_numpy(beta: float) -> None:
    from omnibias.partition.torch import partition_weights as tw

    params, X = _sample()
    ref = partition_weights(params, X, beta)
    got = tw(params, X, beta).detach().cpu().numpy()
    assert np.max(np.abs(got - ref)) < 1e-9

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

Regression example · jax matches numpy

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 -k test_jax_matches_numpy -q

Read the implementation

def partition_weights(params: PartitionParams, X: FloatArray, beta: float) -> FloatArray:
    r"""Soft partition-of-unity weights ``P`` of shape ``(n, 2**depth)``.

    Each row is a probability vector over the regions (non-negative, sums to ``1``);
    ``P[i, l]`` is the soft-AND membership of sample ``i`` in region ``l``.
    """
    G = gate_activations(params, X, beta)
    return _weights_from_gates(G, params.depth)

Related APIsomnibias.partition.gate_activations↗

FUNCTION

region_rule

omnibias.partition.region_ruleSource ↗

region_rule(params: PartitionParams, region: int) -> str

The conjunction of gate conditions selecting region (a hardened if-then clause).

Parameters

Name	Type & default	Contract
paramspositional-or-keyword	PartitionParamsRequired	See the operation contract and implementation below.
regionpositional-or-keyword	intRequired	See the operation contract and implementation below.

Returnsstr

Executed regression · hardened rules are readable

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_hardened_rules_are_readable() -> None:
    cfg = PartitionConfig(n_features=3, depth=2, split_kind="axis", seed=1)
    params = init_params(cfg)
    rules = hardened_rules(params)
    assert len(rules) == cfg.depth
    assert all("x[" in r for r in rules)  # single-feature threshold form
    clause = region_rule(params, region=cfg.n_regions - 1)
    assert "AND" in clause

python -m pytest "packages/omnibias-partition/tests/test_weights.py::test_hardened_rules_are_readable" -q

Read the implementation

def region_rule(params: PartitionParams, region: int) -> str:
    r"""The conjunction of gate conditions selecting ``region`` (a hardened if-then clause)."""
    codes = region_code_matrix(params.depth)
    parts: list[str] = []
    for j in range(params.depth):
        rule = _gate_rule(params, j)
        parts.append(rule if codes[region, j] > 0.5 else f"NOT ({rule})")
    return " AND ".join(parts)

Related APIsomnibias.partition.region_code_matrix↗

CLASS

Arrangement

omnibias.partition.ArrangementSource ↗

class Arrangement

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.

n affine hyperplanes w_i . x = offset_i in R^D.

Parameters

Name	Type & default	Contract
normalsdataclass field	FloatArrayRequired	See the operation contract and implementation below.
offsetsdataclass field	FloatArrayRequired	See the operation contract and implementation below.

Executed regression · g1 general position cell count

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_g1_general_position_cell_count() -> None:
    rng = np.random.default_rng(0)
    for n, d in ((5, 2), (8, 2), (6, 3), (10, 3), (8, 4)):
        arr = general_position_normals(n, d, rng)
        cells = enumerate_cells_vertices(arr)
        assert len(cells) == max_cells(n, d), (n, d, len(cells), max_cells(n, d))

python -m pytest "packages/omnibias-partition/tests/test_arrangement.py::test_g1_general_position_cell_count" -q

Regression example · g1 degenerate strictly smaller

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_g1_degenerate_strictly_smaller() -> None:
    # Two parallel lines in R^2: fewer cells than the simple maximum.
    arr = Arrangement(
        np.array([[1.0, 0.0], [1.0, 0.0], [0.0, 1.0]], dtype=np.float64),
        np.array([-1.0, 1.0, 0.0], dtype=np.float64),
    )
    cells = enumerate_cells_vertices(arr)
    brute = realized_cells(
        arr,
        np.stack(
            np.meshgrid(np.linspace(-2, 2, 41), np.linspace(-2, 2, 41), indexing="ij"),
            axis=-1,
        ).reshape(-1, 2),
    )
    assert len(cells) == len(brute)
    assert len(cells) < max_cells(3, 2)

python -m pytest packages/omnibias-partition/tests/test_arrangement.py -k test_g1_degenerate_strictly_smaller -q

Fields & methods

normalsfield

normals: FloatArray

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

View source ↗

offsetsfield

offsets: FloatArray

Declared FloatArray field on Arrangement. 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 normals.ndim != 2, raises ValueError('normals must have shape (n, D)').
• When offsets.shape[0] != normals.shape[0], raises ValueError('offsets must have length n').

Executed regression · g1 general position cell count

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_g1_general_position_cell_count() -> None:
    rng = np.random.default_rng(0)
    for n, d in ((5, 2), (8, 2), (6, 3), (10, 3), (8, 4)):
        arr = general_position_normals(n, d, rng)
        cells = enumerate_cells_vertices(arr)
        assert len(cells) == max_cells(n, d), (n, d, len(cells), max_cells(n, d))

python -m pytest "packages/omnibias-partition/tests/test_arrangement.py::test_g1_general_position_cell_count" -q

Method implementation

    def __post_init__(self) -> None:
        normals = np.asarray(self.normals, dtype=np.float64)
        offsets = np.asarray(self.offsets, dtype=np.float64).reshape(-1)
        if normals.ndim != 2:
            raise ValueError("normals must have shape (n, D)")
        if offsets.shape[0] != normals.shape[0]:
            raise ValueError("offsets must have length n")
        object.__setattr__(self, "normals", normals)
        object.__setattr__(self, "offsets", offsets)

View source ↗

nproperty

n(self) -> int

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

Returns: int

Executed regression · g1 general position cell count

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_g1_general_position_cell_count() -> None:
    rng = np.random.default_rng(0)
    for n, d in ((5, 2), (8, 2), (6, 3), (10, 3), (8, 4)):
        arr = general_position_normals(n, d, rng)
        cells = enumerate_cells_vertices(arr)
        assert len(cells) == max_cells(n, d), (n, d, len(cells), max_cells(n, d))

python -m pytest "packages/omnibias-partition/tests/test_arrangement.py::test_g1_general_position_cell_count" -q

Method implementation

    def n(self) -> int:
        return int(self.normals.shape[0])

View source ↗

dimproperty

dim(self) -> int

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

Returns: int

Executed regression · g1 general position cell count

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_g1_general_position_cell_count() -> None:
    rng = np.random.default_rng(0)
    for n, d in ((5, 2), (8, 2), (6, 3), (10, 3), (8, 4)):
        arr = general_position_normals(n, d, rng)
        cells = enumerate_cells_vertices(arr)
        assert len(cells) == max_cells(n, d), (n, d, len(cells), max_cells(n, d))

python -m pytest "packages/omnibias-partition/tests/test_arrangement.py::test_g1_general_position_cell_count" -q

Method implementation

    def dim(self) -> int:
        return int(self.normals.shape[1])

View source ↗

Read the implementation

class Arrangement:
    """``n`` affine hyperplanes ``w_i . x = offset_i`` in ``R^D``."""

    normals: FloatArray
    offsets: FloatArray

    def __post_init__(self) -> None:
        normals = np.asarray(self.normals, dtype=np.float64)
        offsets = np.asarray(self.offsets, dtype=np.float64).reshape(-1)
        if normals.ndim != 2:
            raise ValueError("normals must have shape (n, D)")
        if offsets.shape[0] != normals.shape[0]:
            raise ValueError("offsets must have length n")
        object.__setattr__(self, "normals", normals)
        object.__setattr__(self, "offsets", offsets)

    @property
    def n(self) -> int:
        return int(self.normals.shape[0])

    @property
    def dim(self) -> int:
        return int(self.normals.shape[1])

Other public import paths

• omnibias.partition.Arrangement
• omnibias.partition.arrangement.Arrangement

CLASS

CellGapCertificate

omnibias.partition.CellGapCertificateSource ↗

class CellGapCertificate

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.

Sound bound on |soft - hard|. Temperature collapse, not founding.

Parameters

Name	Type & default	Contract
betadataclass field	floatRequired	See the operation contract and implementation below.
bounddataclass field	floatRequired	See the operation contract and implementation below.
measureddataclass field	floatRequired	See the operation contract and implementation below.

Executed regression · g2 gap soundness

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_g2_gap_soundness() -> None:
    rng = np.random.default_rng(1)
    arr = general_position_normals(5, 2, rng)
    cells = enumerate_cells_vertices(arr)
    signs = cells[0]
    grid = np.stack(
        np.meshgrid(np.linspace(-1.5, 1.5, 25), np.linspace(-1.5, 1.5, 25), indexing="ij"),
        axis=-1,
    ).reshape(-1, 2)
    extra = rng.normal(size=(80, 2))
    x = np.concatenate([grid, extra], axis=0)
    for beta in (1.0, 4.0, 16.0):
        cert = certify_cell_gap(arr, x, signs, beta=beta)
        assert cert.is_sound
        assert cert.bound >= cert.measured - 1e-12

python -m pytest "packages/omnibias-partition/tests/test_arrangement.py::test_g2_gap_soundness" -q

Fields & methods

betafield

beta: float

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

View source ↗

boundfield

bound: float

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

View source ↗

measuredfield

measured: float

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

View source ↗

is_soundproperty

is_sound(self) -> bool

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

Returns: bool

Executed regression · g2 gap soundness

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_g2_gap_soundness() -> None:
    rng = np.random.default_rng(1)
    arr = general_position_normals(5, 2, rng)
    cells = enumerate_cells_vertices(arr)
    signs = cells[0]
    grid = np.stack(
        np.meshgrid(np.linspace(-1.5, 1.5, 25), np.linspace(-1.5, 1.5, 25), indexing="ij"),
        axis=-1,
    ).reshape(-1, 2)
    extra = rng.normal(size=(80, 2))
    x = np.concatenate([grid, extra], axis=0)
    for beta in (1.0, 4.0, 16.0):
        cert = certify_cell_gap(arr, x, signs, beta=beta)
        assert cert.is_sound
        assert cert.bound >= cert.measured - 1e-12

python -m pytest "packages/omnibias-partition/tests/test_arrangement.py::test_g2_gap_soundness" -q

Method implementation

    def is_sound(self) -> bool:
        return self.bound + 1e-12 >= self.measured

View source ↗

Read the implementation

class CellGapCertificate:
    """Sound bound on ``|soft - hard|``. Temperature collapse, not founding."""

    beta: float
    bound: float
    measured: float

    @property
    def is_sound(self) -> bool:
        return self.bound + 1e-12 >= self.measured

Other public import paths

• omnibias.partition.CellGapCertificate
• omnibias.partition.arrangement.CellGapCertificate

FUNCTION

certify_cell_gap

omnibias.partition.certify_cell_gapSource ↗

certify_cell_gap(arr: Arrangement, x: FloatArray, signs: Sequence[int], *, beta: float) -> CellGapCertificate

Bound |soft - 1[cell]| from the margin via n (1 - sigmoid(beta m)).

Each factor is in (0,1); the product is at least sigmoid(beta m)^n on the true cell and at most 1 - sigmoid(beta m) off it (one opposing gate). The bound is n (1 - sigmoid_iv(beta * m_lo).lo), which only widens.

Parameters

Name	Type & default	Contract
arrpositional-or-keyword	ArrangementRequired	See the operation contract and implementation below.
xpositional-or-keyword	FloatArrayRequired	See the operation contract and implementation below.
signspositional-or-keyword	Sequence[int]Required	See the operation contract and implementation below.
betakeyword-only	floatRequired	See the operation contract and implementation below.

ReturnsCellGapCertificate

Executed regression · g2 gap soundness

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_g2_gap_soundness() -> None:
    rng = np.random.default_rng(1)
    arr = general_position_normals(5, 2, rng)
    cells = enumerate_cells_vertices(arr)
    signs = cells[0]
    grid = np.stack(
        np.meshgrid(np.linspace(-1.5, 1.5, 25), np.linspace(-1.5, 1.5, 25), indexing="ij"),
        axis=-1,
    ).reshape(-1, 2)
    extra = rng.normal(size=(80, 2))
    x = np.concatenate([grid, extra], axis=0)
    for beta in (1.0, 4.0, 16.0):
        cert = certify_cell_gap(arr, x, signs, beta=beta)
        assert cert.is_sound
        assert cert.bound >= cert.measured - 1e-12

python -m pytest "packages/omnibias-partition/tests/test_arrangement.py::test_g2_gap_soundness" -q

Read the implementation

def certify_cell_gap(
    arr: Arrangement,
    x: FloatArray,
    signs: Sequence[int],
    *,
    beta: float,
) -> CellGapCertificate:
    """Bound ``|soft - 1[cell]|`` from the margin via ``n (1 - sigmoid(beta m))``.

    Each factor is in ``(0,1)``; the product is at least
    ``sigmoid(beta m)^n`` on the true cell and at most
    ``1 - sigmoid(beta m)`` off it (one opposing gate). The bound is
    ``n (1 - sigmoid_iv(beta * m_lo).lo)``, which only widens.
    """
    xv = np.asarray(x, dtype=np.float64)
    if xv.ndim == 1:
        xv = xv.reshape(1, -1)
    z = affine_values(arr, xv)
    s = np.asarray(signs, dtype=np.float64).reshape(1, -1)
    oriented = s * z
    hard = (oriented > 0.0).all(axis=-1).astype(np.float64)
    soft = np.prod(_sigmoid(beta * oriented), axis=-1)
    measured = float(np.max(np.abs(soft - hard)))
    m = np.min(np.abs(z), axis=-1)
    m_lo = float(np.min(m))
    # 1 - sigmoid(beta m) <= 1 / (1 + exp(beta m)); enclose with sigmoid_iv.
    one_minus = Interval.point(1.0) - sigmoid_iv(Interval.point(beta * m_lo))
    bound_iv = Interval.point(float(arr.n)) * one_minus
    return CellGapCertificate(beta=beta, bound=float(bound_iv.hi), measured=measured)

Related APIsomnibias.core.verified.sigmoid_iv↗omnibias.partition.CellGapCertificate↗omnibias.partition.arrangement.affine_values↗

Other public import paths

• omnibias.partition.certify_cell_gap
• omnibias.partition.arrangement.certify_cell_gap

FUNCTION

max_cells

omnibias.partition.max_cellsSource ↗

max_cells(n: int, dim: int) -> int

Maximum number of cells of a simple arrangement of n hyperplanes in R^dim.

Parameters

Name	Type & default	Contract
npositional-or-keyword	intRequired	See the operation contract and implementation below.
dimpositional-or-keyword	intRequired	See the operation contract and implementation below.

Returnsint

Explicit validation & failure cases

• When n < 0 or dim < 0, raises ValueError('n and dim must be non-negative').

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

Executed regression · g1 general position cell count

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_g1_general_position_cell_count() -> None:
    rng = np.random.default_rng(0)
    for n, d in ((5, 2), (8, 2), (6, 3), (10, 3), (8, 4)):
        arr = general_position_normals(n, d, rng)
        cells = enumerate_cells_vertices(arr)
        assert len(cells) == max_cells(n, d), (n, d, len(cells), max_cells(n, d))

python -m pytest "packages/omnibias-partition/tests/test_arrangement.py::test_g1_general_position_cell_count" -q

Regression example · g1 degenerate strictly smaller

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_g1_degenerate_strictly_smaller() -> None:
    # Two parallel lines in R^2: fewer cells than the simple maximum.
    arr = Arrangement(
        np.array([[1.0, 0.0], [1.0, 0.0], [0.0, 1.0]], dtype=np.float64),
        np.array([-1.0, 1.0, 0.0], dtype=np.float64),
    )
    cells = enumerate_cells_vertices(arr)
    brute = realized_cells(
        arr,
        np.stack(
            np.meshgrid(np.linspace(-2, 2, 41), np.linspace(-2, 2, 41), indexing="ij"),
            axis=-1,
        ).reshape(-1, 2),
    )
    assert len(cells) == len(brute)
    assert len(cells) < max_cells(3, 2)

python -m pytest packages/omnibias-partition/tests/test_arrangement.py -k test_g1_degenerate_strictly_smaller -q

Read the implementation

def max_cells(n: int, dim: int) -> int:
    """Maximum number of cells of a simple arrangement of ``n`` hyperplanes in ``R^dim``."""
    if n < 0 or dim < 0:
        raise ValueError("n and dim must be non-negative")
    return int(sum(comb(n, k) for k in range(min(n, dim) + 1)))

Other public import paths

• omnibias.partition.max_cells
• omnibias.partition.arrangement.max_cells

FUNCTION

soft_membership

omnibias.partition.soft_membershipSource ↗

soft_membership(arr: Arrangement, x: FloatArray, signs: Sequence[int], *, beta: float) -> FloatArray

Soft indicator of one cell. beta -> inf is temperature collapse.

Parameters

Name	Type & default	Contract
arrpositional-or-keyword	ArrangementRequired	See the operation contract and implementation below.
xpositional-or-keyword	FloatArrayRequired	See the operation contract and implementation below.
signspositional-or-keyword	Sequence[int]Required	See the operation contract and implementation below.
betakeyword-only	floatRequired	See the operation contract and implementation below.

ReturnsFloatArray

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 · g6 torch runs

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_g6_torch_runs() -> None:
    pytest.importorskip("torch")
    import torch
    from omnibias.graph.arrangement.torch import FaceNet

    torch.set_default_dtype(torch.float64)
    arr = _square()
    samples = np.array([[0.0, 0.0], [2.0, 0.0], [0.0, 2.0], [-2.0, 0.0], [0.0, -2.0]])
    graph = build_arrangement_graph(arr, samples)
    net = FaceNet(2, hidden=4, rounds=2, dtype=torch.float64)
    out = net(graph, arr)
    assert out.shape[0] == len(graph.cells)
    assert torch.isfinite(out).all()

python -m pytest "packages/omnibias-graph/tests/test_arrangement.py::test_g6_torch_runs" -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: FloatArray,
    signs: Sequence[int],
    *,
    beta: float,
) -> FloatArray:
    """Soft indicator of one cell. ``beta -> inf`` is temperature collapse."""
    if beta <= 0.0:
        raise ValueError("beta must be > 0 (temperature collapse axis)")
    z = affine_values(arr, x)
    s = np.asarray(signs, dtype=np.float64).reshape(1, -1)
    return np.prod(_sigmoid(beta * s * z), axis=-1)

Related APIsomnibias.partition.arrangement.affine_values↗

Other public import paths

• omnibias.partition.soft_membership
• omnibias.partition.arrangement.soft_membership

← Package overviewRunnable learning examples →