omnibias.core.multi_index

0.5.0rc1Release candidatePython >=3.10Apache-2.0

Coefficient recurrences, activation contracts, jet algebra, and verified arithmetic shared by the numerical backends.

import omnibias.core.multi_index

Module contract & mathematical conventions

Multi-index combinatorics for multivariate Faà di Bruno jets (pure Python).

This module supplies the backend-agnostic bookkeeping for the *multivariate* multi-layer jet primitive, the dim-D generalisation of the directional (1-D path) jets in omnibias.jax.jet / omnibias.torch.jet.

A multivariate jet truncated at total order N over D variables represents a function g around a base point x0 by its Taylor coefficients

g(x_0 + \delta) = \sum_{|\alpha| \le N} c_\alpha\, \delta^\alpha, \qquad c_\alpha = \frac{D^\alpha g(x_0)}{\alpha!},

where \alpha \in \mathbb{N}^D is a *multi-index*, |\alpha| = \sum_i \alpha_i, \alpha! = \prod_i \alpha_i! and \delta^\alpha = \prod_i \delta_i^{\alpha_i}.

The coefficients c_\alpha are stored densely along the leading axis of a backend array, in the canonical order returned by multi_indices. The only combinatorics the backend kernels need are:

multi_indices -- the canonical, deterministic ordering of all multi-indices with |\alpha| \le N (so torch / jax agree row-for-row);
multiply_table -- the truncated-Cauchy-product table that drives multivariate polynomial multiplication (a*b)_\gamma = \sum_{\alpha+\beta=\gamma} a_\alpha b_\beta, the multivariate replacement for the 1-D convolution in the directional kernel;
multi_index_factorial -- \alpha!, converting between Taylor coefficients and raw partial derivatives.

Everything is pure integer combinatorics (no numpy / torch / jax) and cached on (dim, order) so the two backends are bit-identical by construction.

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

MAX_MULTI_INDICES

omnibias.core.multi_index.MAX_MULTI_INDICESSource ↗

MAX_MULTI_INDICES: int = 200000

Default ceiling on how many multi-indices a single request may materialise. The count is comb(dim + order, dim), which grows fast in *both* arguments, so one joint bound on the result size is the honest guard -- capping dim and order separately would still admit a combinatorial explosion. Checked before anything is allocated, since the count itself is cheap. Settable via set_multi_index_budget.

Regression example · bell and multi index orders are bounded

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_bell_and_multi_index_orders_are_bounded() -> None:
    import pytest
    from omnibias.core.bell import (
        MAX_BELL_NUMBER_ORDER,
        MAX_BELL_ORDER,
        bell_complete,
        bell_number,
        bell_partial,
    )
    from omnibias.core.multi_index import MAX_MULTI_INDICES, multi_indices

    # The partition enumerations are exponential in n, so they cap much lower
    # than the quadratic Bell triangle.
    with pytest.raises(ValueError, match="MAX_BELL_ORDER"):
        bell_complete(MAX_BELL_ORDER + 1)
    with pytest.raises(ValueError, match="MAX_BELL_ORDER"):
        bell_partial(MAX_BELL_ORDER + 1, 2)
    assert bell_number(500) > 0  # the documented in-tree usage still works
    with pytest.raises(ValueError, match="<="):
        bell_number(MAX_BELL_NUMBER_ORDER + 1)

    # One joint bound on the result size, since the count grows in both args.
    with pytest.raises(ValueError, match="multi-index budget"):
        multi_indices(6, 60)
    assert len(multi_indices(3, 4)) == 35
    assert MAX_MULTI_INDICES > 0

python -m pytest packages/omnibias-core/tests/test_polynomials_invariants.py -k test_bell_and_multi_index_orders_are_bounded -q

Read the implementation

MAX_MULTI_INDICES: int = 200_000

CONSTANT

MultiIndex

omnibias.core.multi_index.MultiIndexSource ↗

MultiIndex = tuple[int, ...]

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

MultiIndex = tuple[int, ...]

FUNCTION

get_multi_index_budget

omnibias.core.multi_index.get_multi_index_budgetSource ↗

get_multi_index_budget() -> int

Current process-wide multi-index materialisation ceiling.

Returnsint

Read the implementation

def get_multi_index_budget() -> int:
    """Current process-wide multi-index materialisation ceiling."""
    return _multi_index_budget

FUNCTION

index_position

omnibias.core.index_positionSource ↗

index_position(dim: int, order: int) -> dict[MultiIndex, int]

Map each multi-index to its row in multi_indices.

Parameters

Name	Type & default	Contract
`dim`positional-or-keyword	`int`Required	See the operation contract and implementation below.
`order`positional-or-keyword	`int`Required	See the operation contract and implementation below.

Returnsdict[MultiIndex, int]

Executed regression · value equals jet row0

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_value_equals_jet_row0() -> None:
    net = make_jet_mlp(2, 16, 1, depth=3, seed=1)
    x = _x(2)
    jet = net.jet(x, 2)
    assert jnp.allclose(jet[:, 0, :], net.value(x), atol=1e-14)

python -m pytest "packages/omnibias-jax/tests/test_jet_pinn.py::test_value_equals_jet_row0" -q

Regression example · index position is inverse

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_index_position_is_inverse(dim: int, order: int) -> None:
    idx = multi_indices(dim, order)
    pos = index_position(dim, order)
    for i, alpha in enumerate(idx):
        assert pos[alpha] == i

python -m pytest packages/omnibias-core/tests/test_multi_index.py -k test_index_position_is_inverse -q

Read the implementation

def index_position(dim: int, order: int) -> dict[MultiIndex, int]:
    """Map each multi-index to its row in :func:`multi_indices`."""
    return dict(_index_position(dim, order))

Other public import paths

omnibias.core.index_position
omnibias.core.multi_index.index_position

FUNCTION

multi_index_factorial

omnibias.core.multi_index_factorialSource ↗

multi_index_factorial(alpha: MultiIndex) -> int

Multi-index factorial \alpha! = \prod_i \alpha_i!.

Parameters

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

Returnsint

Explicit validation & failure cases

When a < 0, raises ValueError(f'multi-index entries must be >= 0, got {alpha}').

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

Executed regression · multi index factorial rejects negative

Full file, imports & fixtures ↗

def test_multi_index_factorial_rejects_negative() -> None:
    with pytest.raises(ValueError, match=">= 0"):
        multi_index_factorial((1, -1))

python -m pytest "packages/omnibias-core/tests/test_multi_index.py::test_multi_index_factorial_rejects_negative" -q

Regression example · multi index factorial

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_multi_index_factorial() -> None:
    assert multi_index_factorial((0, 0)) == 1
    assert multi_index_factorial((1, 0)) == 1
    assert multi_index_factorial((2, 3)) == factorial(2) * factorial(3)
    assert multi_index_factorial((4,)) == 24

python -m pytest packages/omnibias-core/tests/test_multi_index.py -k test_multi_index_factorial -q

Read the implementation

def multi_index_factorial(alpha: MultiIndex) -> int:
    """Multi-index factorial :math:`\\alpha! = \\prod_i \\alpha_i!`."""
    result = 1
    for a in alpha:
        if a < 0:
            raise ValueError(f"multi-index entries must be >= 0, got {alpha}")
        result *= factorial(a)
    return result

Other public import paths

omnibias.core.multi_index_factorial
omnibias.core.multi_index.multi_index_factorial

FUNCTION

multi_indices

omnibias.core.multi_indicesSource ↗

multi_indices(dim: int, order: int) -> list[MultiIndex]

Canonical ordering of all multi-indices with |alpha| <= order.

Indices are sorted by total degree |alpha| and then lexicographically, so element 0 is always the zero multi-index (0, ..., 0) and the next dim entries are the unit vectors e_0, ..., e_{dim-1} (the gradient block). The ordering is deterministic, so dense jets share row indexing across backends.

Parameters

Name	Type & default	Contract
`dim`positional-or-keyword	`int`Required	See the operation contract and implementation below.
`order`positional-or-keyword	`int`Required	See the operation contract and implementation below.

Returnslist[MultiIndex]

Executed regression · first row is zero index

Full file, imports & fixtures ↗

def test_first_row_is_zero_index() -> None:
    assert multi_indices(3, 2)[0] == (0, 0, 0)

python -m pytest "packages/omnibias-core/tests/test_multi_index.py::test_first_row_is_zero_index" -q

Regression example · multi indices rejects bad dim

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_multi_indices_rejects_bad_dim() -> None:
    with pytest.raises(ValueError, match="dim must be >= 1"):
        multi_indices(0, 2)

python -m pytest packages/omnibias-core/tests/test_multi_index.py -k test_multi_indices_rejects_bad_dim -q

Read the implementation

def multi_indices(dim: int, order: int) -> list[MultiIndex]:
    """Canonical ordering of all multi-indices with ``|alpha| <= order``.

    Indices are sorted by total degree ``|alpha|`` and then lexicographically,
    so element ``0`` is always the zero multi-index ``(0, ..., 0)`` and the next
    ``dim`` entries are the unit vectors ``e_0, ..., e_{dim-1}`` (the gradient
    block). The ordering is deterministic, so dense jets share row indexing
    across backends.
    """
    return list(_multi_indices(dim, order))

Other public import paths

omnibias.core.multi_indices
omnibias.core.multi_index.multi_indices

FUNCTION

multiply_table

omnibias.core.multiply_tableSource ↗

multiply_table(dim: int, order: int) -> list[tuple[int, int, int]]

Truncated Cauchy-product table for multivariate polynomial multiply.

Returns a list of (gamma, alpha, beta) *row-index* triples such that indices[alpha] + indices[beta] == indices[gamma] componentwise (with indices = multi_indices(dim, order)). The product of two dense jets a, b is then c[gamma] = sum a[alpha] * b[beta] over all triples sharing gamma; pairs whose sum exceeds order are absent (truncated).

Parameters

Name	Type & default	Contract
`dim`positional-or-keyword	`int`Required	See the operation contract and implementation below.
`order`positional-or-keyword	`int`Required	See the operation contract and implementation below.

Returnslist[tuple[int, int, int]]

Executed regression · value equals jet row0

Full file, imports & fixtures ↗

def test_value_equals_jet_row0() -> None:
    net = make_jet_mlp(2, 16, 1, depth=3, seed=1)
    x = _x(2)
    jet = net.jet(x, 2)
    assert jnp.allclose(jet[:, 0, :], net.value(x), atol=1e-14)

python -m pytest "packages/omnibias-jax/tests/test_jet_pinn.py::test_value_equals_jet_row0" -q

Regression example · multiply table consistency

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_multiply_table_consistency(dim: int, order: int) -> None:
    idx = multi_indices(dim, order)
    for g, a, b in multiply_table(dim, order):
        gamma, alpha, beta = idx[g], idx[a], idx[b]
        assert tuple(alpha[i] + beta[i] for i in range(dim)) == gamma

python -m pytest packages/omnibias-core/tests/test_multi_index.py -k test_multiply_table_consistency -q

Read the implementation

def multiply_table(dim: int, order: int) -> list[tuple[int, int, int]]:
    """Truncated Cauchy-product table for multivariate polynomial multiply.

    Returns a list of ``(gamma, alpha, beta)`` *row-index* triples such that
    ``indices[alpha] + indices[beta] == indices[gamma]`` componentwise (with
    ``indices = multi_indices(dim, order)``). The product of two dense jets
    ``a``, ``b`` is then ``c[gamma] = sum a[alpha] * b[beta]`` over all triples
    sharing ``gamma``; pairs whose sum exceeds ``order`` are absent (truncated).
    """
    return list(_multiply_table(dim, order))

Other public import paths

omnibias.core.multiply_table
omnibias.core.multi_index.multiply_table

FUNCTION

num_multi_indices

omnibias.core.num_multi_indicesSource ↗

num_multi_indices(dim: int, order: int) -> int

Number of multi-indices with |alpha| <= order in dim variables.

Equals \binom{D + N}{D}, the dimension of the truncated Taylor coefficient space (and the length of multi_indices).

Parameters

Name	Type & default	Contract
`dim`positional-or-keyword	`int`Required	See the operation contract and implementation below.
`order`positional-or-keyword	`int`Required	See the operation contract and implementation below.

Returnsint

Explicit validation & failure cases

When dim < 1, raises ValueError(f'dim must be >= 1, got {dim}').
When order < 0, raises ValueError(f'order must be >= 0, got {order}').

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

Executed regression · coefficient count must match

Full file, imports & fixtures ↗

def test_coefficient_count_must_match() -> None:
    with pytest.raises(ValueError):
        TaylorModelMV((0.0, 0.0), (1.0, 1.0), 2, [Interval.point(0.0)], Interval.point(0.0))

python -m pytest "packages/omnibias-core/tests/verified/test_taylor_model_mv.py::test_coefficient_count_must_match" -q

Regression example · count matches binomial

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_count_matches_binomial(dim: int, order: int) -> None:
    idx = multi_indices(dim, order)
    assert len(idx) == comb(dim + order, dim)
    assert len(idx) == num_multi_indices(dim, order)

python -m pytest packages/omnibias-core/tests/test_multi_index.py -k test_count_matches_binomial -q

Read the implementation

def num_multi_indices(dim: int, order: int) -> int:
    """Number of multi-indices with ``|alpha| <= order`` in ``dim`` variables.

    Equals :math:`\\binom{D + N}{D}`, the dimension of the truncated Taylor
    coefficient space (and the length of :func:`multi_indices`).
    """
    if dim < 1:
        raise ValueError(f"dim must be >= 1, got {dim}")
    if order < 0:
        raise ValueError(f"order must be >= 0, got {order}")
    return comb(dim + order, dim)

Other public import paths

omnibias.core.num_multi_indices
omnibias.core.multi_index.num_multi_indices

FUNCTION

set_multi_index_budget

omnibias.core.multi_index.set_multi_index_budgetSource ↗

set_multi_index_budget(budget: int | None) -> None

Override the process-wide multi-index materialisation ceiling.

Pass None to restore the default (MAX_MULTI_INDICES). Raising this ceiling is a real memory hazard: comb(dim + order, dim) grows combinatorially in both arguments. For a deep network's Laplacian or poly-Laplacian, prefer the ceiling-free contraction kernels (deep_field_laplacian / deep_field_polylaplacian in omnibias.jax.laplacian / omnibias.torch.laplacian) over raising this budget -- they never materialise the full multivariate jet.

Parameters

Name	Type & default	Contract
`budget`positional-or-keyword	`int \| None`Required	See the operation contract and implementation below.

ReturnsNone

Explicit validation & failure cases

When budget < 1, raises ValueError(f'budget must be >= 1, got {budget}').

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

Read the implementation

def set_multi_index_budget(budget: int | None) -> None:
    """Override the process-wide multi-index materialisation ceiling.

    Pass ``None`` to restore the default (:data:`MAX_MULTI_INDICES`).
    Raising this ceiling is a real memory hazard: ``comb(dim + order, dim)``
    grows combinatorially in both arguments. For a deep network's Laplacian
    or poly-Laplacian, prefer the ceiling-free contraction kernels
    (``deep_field_laplacian`` / ``deep_field_polylaplacian`` in
    :mod:`omnibias.jax.laplacian` / :mod:`omnibias.torch.laplacian`) over
    raising this budget -- they never materialise the full multivariate jet.
    """
    global _multi_index_budget
    if budget is None:
        _multi_index_budget = MAX_MULTI_INDICES
        return
    if budget < 1:
        raise ValueError(f"budget must be >= 1, got {budget}")
    _multi_index_budget = int(budget)

← Package overview Runnable learning examples →

omnibias.core.multi_index

MAX_MULTI_INDICES​

MultiIndex​

get_multi_index_budget​

index_position​

Parameters

multi_index_factorial​

Parameters

multi_indices​

Parameters

multiply_table​

Parameters

num_multi_indices​

Parameters

set_multi_index_budget​

Parameters

MAX_MULTI_INDICES

MultiIndex

get_multi_index_budget

index_position

multi_index_factorial

multi_indices

multiply_table

num_multi_indices

set_multi_index_budget