API REFERENCE/omnibias-jax

omnibias.jax.refine

0.5.0rc1Release candidatePython >=3.10Apache-2.0

JAX activation derivatives, functional network jets, direct contractions, and parameter-differentiable numerical models.

import omnibias.jax.refine

Module contract & mathematical conventions

Adaptive pack refinement (theory 03-13), JAX twin.

Bit-identical partner of omnibias.torch.refine on the same pack list and the same residual / jet numbers. Structural decisions run through omnibias.core.refine.refine_bank (host-side). Do not jit the driver.

There is no jax GrowableOperatorMultiBiasUnit. Literal-K growth stays on the torch growable unit; this module's *p*-type move is the collapsed sibling (new sigma^(n+1) slot with c = 0).

Enable 64-bit JAX before the first array if you need torch parity (omnibias.jax.precision).

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 →

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
Indicator	omnibias.core.Indicator
RefinePolicy	omnibias.core.RefinePolicy
RefineReport	omnibias.core.RefineReport
RefinedPack	omnibias.core.RefinedPack

CLASS

AdaptivePackBank

omnibias.jax.AdaptivePackBankSource ↗

class AdaptivePackBank

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.

Functional tempered pack bank (pytree of slot arrays).

Parameters

Name	Type & default	Contract
centersdataclass field	ArrayRequired	See the operation contract and implementation below.
weightsdataclass field	ArrayRequired	See the operation contract and implementation below.
scalesdataclass field	ArrayRequired	See the operation contract and implementation below.
ordersdataclass field	ArrayRequired	See the operation contract and implementation below.
agesdataclass field	ArrayRequired	See the operation contract and implementation below.
activedataclass field	ArrayRequired	See the operation contract and implementation below.
max_orderdataclass field	intRequired	See the operation contract and implementation below.
basedataclass field	strRequired	See the operation contract and implementation below.

Executed regression · g3 scale flow

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_g3_scale_flow() -> None:
    bank = _initial_bank()
    probes = jnp.linspace(0.0, 1.0, 65)

    def residual(x: Array) -> Array:
        return jnp.exp(-100.0 * x)

    _bank, report = refine(
        bank,
        residual,
        probes,
        RefinePolicy(
            indicator=Indicator.SCALE_FLOW,
            birth_threshold=0.1,
            death_threshold=0.0,
            min_age=10_000,
            hysteresis=1.0,
            min_scale_ratio=1.0,
        ),
        step=0,
        probe_jet=(_BL_CENTER, _BL_DERIVS),
    )
    assert report.proposed_scale == pytest.approx(_BL_SCALE)

python -m pytest "packages/omnibias-jax/tests/test_refine.py::test_g3_scale_flow" -q

Fields & methods

centersfield

centers: Array

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

View source ↗

weightsfield

weights: Array

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

View source ↗

scalesfield

scales: Array

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

View source ↗

ordersfield

orders: Array

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

View source ↗

agesfield

ages: Array

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

View source ↗

activefield

active: Array

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

View source ↗

max_orderfield

max_order: int

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

View source ↗

basefield

base: str

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

View source ↗

max_packsproperty

max_packs(self) -> int

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

Returns: int

Executed regression · g3 scale flow

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_g3_scale_flow() -> None:
    bank = _initial_bank()
    probes = jnp.linspace(0.0, 1.0, 65)

    def residual(x: Array) -> Array:
        return jnp.exp(-100.0 * x)

    _bank, report = refine(
        bank,
        residual,
        probes,
        RefinePolicy(
            indicator=Indicator.SCALE_FLOW,
            birth_threshold=0.1,
            death_threshold=0.0,
            min_age=10_000,
            hysteresis=1.0,
            min_scale_ratio=1.0,
        ),
        step=0,
        probe_jet=(_BL_CENTER, _BL_DERIVS),
    )
    assert report.proposed_scale == pytest.approx(_BL_SCALE)

python -m pytest "packages/omnibias-jax/tests/test_refine.py::test_g3_scale_flow" -q

Method implementation

    def max_packs(self) -> int:
        return int(self.centers.shape[0])

View source ↗

n_activeproperty

n_active(self) -> int

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

Returns: int

Method implementation

    def n_active(self) -> int:
        return int(jnp.asarray(self.active).sum())

View source ↗

Read the implementation

class AdaptivePackBank:
    """Functional tempered pack bank (pytree of slot arrays)."""

    centers: Array
    weights: Array
    scales: Array
    orders: Array
    ages: Array
    active: Array
    max_order: int
    base: str

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

    @property
    def n_active(self) -> int:
        return int(jnp.asarray(self.active).sum())

Other public import paths

• omnibias.jax.AdaptivePackBank
• omnibias.jax.refine.AdaptivePackBank

FUNCTION

bank_active_packs

omnibias.jax.refine.bank_active_packsSource ↗

bank_active_packs(bank: AdaptivePackBank) -> list[RefinedPack]

The module contract describes the context for this operation. Its signature, explicit validation, return expressions, and implementation are recorded below.

Parameters

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

Returnslist[RefinedPack]

Executed regression · g3 scale flow

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_g3_scale_flow() -> None:
    bank = _initial_bank()
    probes = jnp.linspace(0.0, 1.0, 65)

    def residual(x: Array) -> Array:
        return jnp.exp(-100.0 * x)

    _bank, report = refine(
        bank,
        residual,
        probes,
        RefinePolicy(
            indicator=Indicator.SCALE_FLOW,
            birth_threshold=0.1,
            death_threshold=0.0,
            min_age=10_000,
            hysteresis=1.0,
            min_scale_ratio=1.0,
        ),
        step=0,
        probe_jet=(_BL_CENTER, _BL_DERIVS),
    )
    assert report.proposed_scale == pytest.approx(_BL_SCALE)

python -m pytest "packages/omnibias-jax/tests/test_refine.py::test_g3_scale_flow" -q

Read the implementation

def bank_active_packs(bank: AdaptivePackBank) -> list[RefinedPack]:
    out: list[RefinedPack] = []
    for i in range(bank.max_packs):
        if bool(bank.active[i]):
            out.append(
                RefinedPack(
                    order=int(bank.orders[i]),
                    center=float(bank.centers[i]),
                    weight=float(bank.weights[i]),
                    scale=float(bank.scales[i]),
                    age=int(bank.ages[i]),
                )
            )
    return out

Related APIsomnibias.core.RefinedPack↗

FUNCTION

bank_forward

omnibias.jax.bank_forwardSource ↗

bank_forward(bank: AdaptivePackBank, z: Array) -> Array

1-D field z → same leading shape.

Parameters

Name	Type & default	Contract
bankpositional-or-keyword	AdaptivePackBankRequired	See the operation contract and implementation below.
zpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.

ReturnsArray

Executed regression · g3 scale flow

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_g3_scale_flow() -> None:
    bank = _initial_bank()
    probes = jnp.linspace(0.0, 1.0, 65)

    def residual(x: Array) -> Array:
        return jnp.exp(-100.0 * x)

    _bank, report = refine(
        bank,
        residual,
        probes,
        RefinePolicy(
            indicator=Indicator.SCALE_FLOW,
            birth_threshold=0.1,
            death_threshold=0.0,
            min_age=10_000,
            hysteresis=1.0,
            min_scale_ratio=1.0,
        ),
        step=0,
        probe_jet=(_BL_CENTER, _BL_DERIVS),
    )
    assert report.proposed_scale == pytest.approx(_BL_SCALE)

python -m pytest "packages/omnibias-jax/tests/test_refine.py::test_g3_scale_flow" -q

Regression example · birth bit identical

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_birth_bit_identical() -> None:
    bank = _initial_bank()
    probes = jnp.linspace(0.0, 1.0, 33)

    def residual(x: Array) -> Array:
        return jnp.ones_like(x)

    before = bank_forward(bank, probes)
    bank, report = refine(
        bank,
        residual,
        probes,
        RefinePolicy(
            indicator=Indicator.SINGULARITY,
            birth_threshold=0.1,
            death_threshold=0.0,
            min_age=10_000,
            hysteresis=1.0,
            min_scale_ratio=1.0,
            debug_zero_perturbation=True,
        ),
        step=0,
        probe_jet=(_BL_CENTER, _BL_DERIVS),
    )
    after = bank_forward(bank, probes)
    assert report.born
    assert bool(jnp.array_equal(before, after))
    assert report.proposed_scale == pytest.approx(_BL_SCALE)

python -m pytest packages/omnibias-jax/tests/test_refine.py -k test_birth_bit_identical -q

Read the implementation

def bank_forward(bank: AdaptivePackBank, z: Array) -> Array:
    """1-D field ``z`` → same leading shape."""
    spec = get_activation(bank.base)
    flat = _as_1d(z)
    out = jnp.zeros_like(flat)
    for i in range(bank.max_packs):
        if not bool(bank.active[i]):
            continue
        n = int(bank.orders[i])
        out = out + pack_response(
            flat,
            bank.centers[i],
            bank.weights[i],
            bank.scales[i],
            n,
            spec,
        )
    if z.ndim != 1 and tuple(z.shape[-1:]) == (1,):
        return out.reshape(z.shape)
    return out

Related APIsomnibias.jax.get_activation↗omnibias.jax.refine.pack_response↗

Other public import paths

• omnibias.jax.bank_forward
• omnibias.jax.refine.bank_forward

FUNCTION

bank_sync

omnibias.jax.refine.bank_syncSource ↗

bank_sync(bank: AdaptivePackBank, packs: Sequence[RefinedPack]) -> AdaptivePackBank

Write a refined list back onto slots, keeping live packs in place.

Parameters

Name	Type & default	Contract
bankpositional-or-keyword	AdaptivePackBankRequired	See the operation contract and implementation below.
packspositional-or-keyword	Sequence[RefinedPack]Required	See the operation contract and implementation below.

ReturnsAdaptivePackBank

Explicit validation & failure cases

• When free is None, raises RuntimeError(f'max_packs={bank.max_packs} exhausted').

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

Executed regression · g3 scale flow

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_g3_scale_flow() -> None:
    bank = _initial_bank()
    probes = jnp.linspace(0.0, 1.0, 65)

    def residual(x: Array) -> Array:
        return jnp.exp(-100.0 * x)

    _bank, report = refine(
        bank,
        residual,
        probes,
        RefinePolicy(
            indicator=Indicator.SCALE_FLOW,
            birth_threshold=0.1,
            death_threshold=0.0,
            min_age=10_000,
            hysteresis=1.0,
            min_scale_ratio=1.0,
        ),
        step=0,
        probe_jet=(_BL_CENTER, _BL_DERIVS),
    )
    assert report.proposed_scale == pytest.approx(_BL_SCALE)

python -m pytest "packages/omnibias-jax/tests/test_refine.py::test_g3_scale_flow" -q

Read the implementation

def bank_sync(bank: AdaptivePackBank, packs: Sequence[RefinedPack]) -> AdaptivePackBank:
    """Write a refined list back onto slots, keeping live packs in place."""
    occupied: dict[tuple[int, float, float, float], int] = {}
    for i in range(bank.max_packs):
        if bool(bank.active[i]):
            occupied[_key(_pack_at(bank, i))] = i
    centers = [float(v) for v in bank.centers]
    weights = [float(v) for v in bank.weights]
    scales = [float(v) for v in bank.scales]
    orders = [int(v) for v in bank.orders]
    ages = [int(v) for v in bank.ages]
    active = [bool(v) for v in bank.active]
    assigned: set[int] = set()
    for pack in packs:
        key = _key(pack)
        slot = occupied.get(key)
        if slot is not None and slot not in assigned:
            ages[slot] = pack.age
            assigned.add(slot)
            continue
        free = next((i for i in range(bank.max_packs) if i not in assigned), None)
        if free is None:
            raise RuntimeError(f"max_packs={bank.max_packs} exhausted")
        centers[free] = pack.center
        weights[free] = pack.weight
        scales[free] = pack.scale
        orders[free] = pack.order
        ages[free] = pack.age
        active[free] = True
        assigned.add(free)
    for i in range(bank.max_packs):
        if i not in assigned:
            active[i] = False
            weights[i] = 0.0
            ages[i] = 0
    return replace(
        bank,
        centers=jnp.asarray(centers, dtype=bank.centers.dtype),
        weights=jnp.asarray(weights, dtype=bank.weights.dtype),
        scales=jnp.asarray(scales, dtype=bank.scales.dtype),
        orders=jnp.asarray(orders, dtype=bank.orders.dtype),
        ages=jnp.asarray(ages, dtype=bank.ages.dtype),
        active=jnp.asarray(active, dtype=bool),
    )

FUNCTION

init_pack_bank

omnibias.jax.init_pack_bankSource ↗

init_pack_bank(packs: Sequence[RefinedPack], *, max_packs: int=16, max_order: int=8, base: str | JaxActivationSpec='exp') -> AdaptivePackBank

Allocate slots and write the initial live packs into the prefix.

Parameters

Name	Type & default	Contract
packspositional-or-keyword	Sequence[RefinedPack]Required	See the operation contract and implementation below.
max_packskeyword-only	intDefault: 16	See the operation contract and implementation below.
max_orderkeyword-only	intDefault: 8	See the operation contract and implementation below.
basekeyword-only	str | JaxActivationSpecDefault: 'exp'	See the operation contract and implementation below.

ReturnsAdaptivePackBank

Explicit validation & failure cases

• When int(max_packs) < 1, raises ValueError(f'max_packs must be >= 1, got {max_packs}').
• When int(max_order) < 0, raises ValueError(f'max_order must be >= 0, got {max_order}').
• When len(live) > int(max_packs), raises ValueError(f'{len(live)} initial packs exceed max_packs={max_packs}').
• When act.fastpath is None, raises NotImplementedError(f'Activation {act.name!r} has no closed-form derivative kernel').
• When pack.order > int(max_order), raises ValueError(f'pack order {pack.order} exceeds max_order={max_order}').

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

Executed regression · g3 scale flow

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_g3_scale_flow() -> None:
    bank = _initial_bank()
    probes = jnp.linspace(0.0, 1.0, 65)

    def residual(x: Array) -> Array:
        return jnp.exp(-100.0 * x)

    _bank, report = refine(
        bank,
        residual,
        probes,
        RefinePolicy(
            indicator=Indicator.SCALE_FLOW,
            birth_threshold=0.1,
            death_threshold=0.0,
            min_age=10_000,
            hysteresis=1.0,
            min_scale_ratio=1.0,
        ),
        step=0,
        probe_jet=(_BL_CENTER, _BL_DERIVS),
    )
    assert report.proposed_scale == pytest.approx(_BL_SCALE)

python -m pytest "packages/omnibias-jax/tests/test_refine.py::test_g3_scale_flow" -q

Regression example · birth bit identical

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_birth_bit_identical() -> None:
    bank = _initial_bank()
    probes = jnp.linspace(0.0, 1.0, 33)

    def residual(x: Array) -> Array:
        return jnp.ones_like(x)

    before = bank_forward(bank, probes)
    bank, report = refine(
        bank,
        residual,
        probes,
        RefinePolicy(
            indicator=Indicator.SINGULARITY,
            birth_threshold=0.1,
            death_threshold=0.0,
            min_age=10_000,
            hysteresis=1.0,
            min_scale_ratio=1.0,
            debug_zero_perturbation=True,
        ),
        step=0,
        probe_jet=(_BL_CENTER, _BL_DERIVS),
    )
    after = bank_forward(bank, probes)
    assert report.born
    assert bool(jnp.array_equal(before, after))
    assert report.proposed_scale == pytest.approx(_BL_SCALE)

python -m pytest packages/omnibias-jax/tests/test_refine.py -k test_birth_bit_identical -q

Read the implementation

def init_pack_bank(
    packs: Sequence[RefinedPack],
    *,
    max_packs: int = 16,
    max_order: int = 8,
    base: str | JaxActivationSpec = "exp",
) -> AdaptivePackBank:
    """Allocate slots and write the initial live packs into the prefix."""
    if int(max_packs) < 1:
        raise ValueError(f"max_packs must be >= 1, got {max_packs}")
    if int(max_order) < 0:
        raise ValueError(f"max_order must be >= 0, got {max_order}")
    live = tuple(packs)
    if len(live) > int(max_packs):
        raise ValueError(f"{len(live)} initial packs exceed max_packs={max_packs}")
    act = get_activation(base) if isinstance(base, str) else base
    if act.fastpath is None:
        raise NotImplementedError(
            f"Activation {act.name!r} has no closed-form derivative kernel"
        )
    for pack in live:
        if pack.order > int(max_order):
            raise ValueError(f"pack order {pack.order} exceeds max_order={max_order}")
        act.fastpath(jnp.zeros(()), pack.order)
    centers = [0.0] * int(max_packs)
    weights = [0.0] * int(max_packs)
    scales = [1.0] * int(max_packs)
    orders = [0] * int(max_packs)
    ages = [0] * int(max_packs)
    active = [False] * int(max_packs)
    for i, pack in enumerate(live):
        centers[i] = pack.center
        weights[i] = pack.weight
        scales[i] = pack.scale
        orders[i] = pack.order
        ages[i] = pack.age
        active[i] = True
    return AdaptivePackBank(
        centers=jnp.asarray(centers),
        weights=jnp.asarray(weights),
        scales=jnp.asarray(scales),
        orders=jnp.asarray(orders, dtype=jnp.int32),
        ages=jnp.asarray(ages, dtype=jnp.int32),
        active=jnp.asarray(active, dtype=bool),
        max_order=int(max_order),
        base=act.name,
    )

Related APIsomnibias.jax.get_activation↗omnibias.jax.AdaptivePackBank↗

Other public import paths

• omnibias.jax.init_pack_bank
• omnibias.jax.refine.init_pack_bank

FUNCTION

pack_response

omnibias.jax.refine.pack_responseSource ↗

pack_response(z: Array, center: Array, weight: Array, scale: Array, order: int, spec: JaxActivationSpec) -> Array

c * alpha^n sigma^(n)(alpha (z - center)).

Parameters

Name	Type & default	Contract
zpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
centerpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
weightpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
scalepositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
orderpositional-or-keyword	intRequired	See the operation contract and implementation below.
specpositional-or-keyword	JaxActivationSpecRequired	See the operation contract and implementation below.

ReturnsArray

Explicit validation & failure cases

• When spec.fastpath is None, raises NotImplementedError(f'Activation {spec.name!r} has no closed-form derivative kernel').
• 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 · g3 scale flow

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_g3_scale_flow() -> None:
    bank = _initial_bank()
    probes = jnp.linspace(0.0, 1.0, 65)

    def residual(x: Array) -> Array:
        return jnp.exp(-100.0 * x)

    _bank, report = refine(
        bank,
        residual,
        probes,
        RefinePolicy(
            indicator=Indicator.SCALE_FLOW,
            birth_threshold=0.1,
            death_threshold=0.0,
            min_age=10_000,
            hysteresis=1.0,
            min_scale_ratio=1.0,
        ),
        step=0,
        probe_jet=(_BL_CENTER, _BL_DERIVS),
    )
    assert report.proposed_scale == pytest.approx(_BL_SCALE)

python -m pytest "packages/omnibias-jax/tests/test_refine.py::test_g3_scale_flow" -q

Read the implementation

def pack_response(
    z: Array,
    center: Array,
    weight: Array,
    scale: Array,
    order: int,
    spec: JaxActivationSpec,
) -> Array:
    """``c * alpha^n sigma^(n)(alpha (z - center))``."""
    if spec.fastpath is None:
        raise NotImplementedError(
            f"Activation {spec.name!r} has no closed-form derivative kernel"
        )
    if order < 0:
        raise ValueError(f"order must be >= 0, got {order}")
    u = scale * (z - center)
    return weight * (scale ** order) * spec.fastpath(u, order)

FUNCTION

pack_term_norms

omnibias.jax.refine.pack_term_normsSource ↗

pack_term_norms(bank: AdaptivePackBank, z: Array) -> list[float]

The module contract describes the context for this operation. Its signature, explicit validation, return expressions, and implementation are recorded below.

Parameters

Name	Type & default	Contract
bankpositional-or-keyword	AdaptivePackBankRequired	See the operation contract and implementation below.
zpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.

Returnslist[float]

Executed regression · g3 scale flow

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_g3_scale_flow() -> None:
    bank = _initial_bank()
    probes = jnp.linspace(0.0, 1.0, 65)

    def residual(x: Array) -> Array:
        return jnp.exp(-100.0 * x)

    _bank, report = refine(
        bank,
        residual,
        probes,
        RefinePolicy(
            indicator=Indicator.SCALE_FLOW,
            birth_threshold=0.1,
            death_threshold=0.0,
            min_age=10_000,
            hysteresis=1.0,
            min_scale_ratio=1.0,
        ),
        step=0,
        probe_jet=(_BL_CENTER, _BL_DERIVS),
    )
    assert report.proposed_scale == pytest.approx(_BL_SCALE)

python -m pytest "packages/omnibias-jax/tests/test_refine.py::test_g3_scale_flow" -q

Read the implementation

def pack_term_norms(bank: AdaptivePackBank, z: Array) -> list[float]:
    spec = get_activation(bank.base)
    flat = _as_1d(z)
    norms: list[float] = []
    for i in range(bank.max_packs):
        if not bool(bank.active[i]):
            continue
        n = int(bank.orders[i])
        term = pack_response(
            flat,
            bank.centers[i],
            bank.weights[i],
            bank.scales[i],
            n,
            spec,
        )
        norms.append(float(jnp.sqrt(jnp.sum(term * term))))
    return norms

Related APIsomnibias.jax.get_activation↗omnibias.jax.refine.pack_response↗

FUNCTION

refine

omnibias.jax.refine.refineSource ↗

refine(bank: AdaptivePackBank, residual_fn: ResidualFn, probes: Array, policy: RefinePolicy, *, step: int, target_fn: ResidualFn | None=None, probe_jet: tuple[float, Sequence[float]] | None=None) -> tuple[AdaptivePackBank, RefineReport]

One refine tick. Returns the updated bank (functional).

Parameters

Name	Type & default	Contract
bankpositional-or-keyword	AdaptivePackBankRequired	See the operation contract and implementation below.
residual_fnpositional-or-keyword	ResidualFnRequired	See the operation contract and implementation below.
probespositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
policypositional-or-keyword	RefinePolicyRequired	See the operation contract and implementation below.
stepkeyword-only	intRequired	See the operation contract and implementation below.
target_fnkeyword-only	ResidualFn | NoneDefault: None	See the operation contract and implementation below.
probe_jetkeyword-only	tuple[float, Sequence[float]] | NoneDefault: None	See the operation contract and implementation below.

Returnstuple[AdaptivePackBank, RefineReport]

Explicit validation & failure cases

• When residual.shape != probes_1d.shape, raises ValueError(f'residual shape {tuple(residual.shape)} != probes {tuple(probes_1d.shape)}').

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

Executed regression · g3 scale flow

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_g3_scale_flow() -> None:
    bank = _initial_bank()
    probes = jnp.linspace(0.0, 1.0, 65)

    def residual(x: Array) -> Array:
        return jnp.exp(-100.0 * x)

    _bank, report = refine(
        bank,
        residual,
        probes,
        RefinePolicy(
            indicator=Indicator.SCALE_FLOW,
            birth_threshold=0.1,
            death_threshold=0.0,
            min_age=10_000,
            hysteresis=1.0,
            min_scale_ratio=1.0,
        ),
        step=0,
        probe_jet=(_BL_CENTER, _BL_DERIVS),
    )
    assert report.proposed_scale == pytest.approx(_BL_SCALE)

python -m pytest "packages/omnibias-jax/tests/test_refine.py::test_g3_scale_flow" -q

Regression example · birth bit identical

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_birth_bit_identical() -> None:
    bank = _initial_bank()
    probes = jnp.linspace(0.0, 1.0, 33)

    def residual(x: Array) -> Array:
        return jnp.ones_like(x)

    before = bank_forward(bank, probes)
    bank, report = refine(
        bank,
        residual,
        probes,
        RefinePolicy(
            indicator=Indicator.SINGULARITY,
            birth_threshold=0.1,
            death_threshold=0.0,
            min_age=10_000,
            hysteresis=1.0,
            min_scale_ratio=1.0,
            debug_zero_perturbation=True,
        ),
        step=0,
        probe_jet=(_BL_CENTER, _BL_DERIVS),
    )
    after = bank_forward(bank, probes)
    assert report.born
    assert bool(jnp.array_equal(before, after))
    assert report.proposed_scale == pytest.approx(_BL_SCALE)

python -m pytest packages/omnibias-jax/tests/test_refine.py -k test_birth_bit_identical -q

Read the implementation

def refine(
    bank: AdaptivePackBank,
    residual_fn: ResidualFn,
    probes: Array,
    policy: RefinePolicy,
    *,
    step: int,
    target_fn: ResidualFn | None = None,
    probe_jet: tuple[float, Sequence[float]] | None = None,
) -> tuple[AdaptivePackBank, RefineReport]:
    """One refine tick. Returns the updated bank (functional)."""
    del step
    probes_1d = _as_1d(probes)
    residual = _as_1d(residual_fn(probes_1d))
    if residual.shape != probes_1d.shape:
        raise ValueError(
            f"residual shape {tuple(residual.shape)} != probes {tuple(probes_1d.shape)}"
        )
    before = None
    if policy.debug_zero_perturbation:
        before = [float(v) for v in bank_forward(bank, probes_1d).reshape(-1)]

    packs = bank_active_packs(bank)
    field = bank_forward(bank, probes_1d)
    field_norm = float(jnp.sqrt(jnp.sum(field * field)))
    contrib = pack_contributions(pack_term_norms(bank, probes_1d), field_norm)
    residuals = [float(v) for v in residual]
    probe_list = [float(v) for v in probes_1d]

    certified: list[float] | None = None
    if policy.indicator is Indicator.CERTIFIED_ERROR:
        if len(probe_list) < 2:
            spacing = 1.0
        else:
            spacing = abs(probe_list[1] - probe_list[0])
        certified = []
        for x in probe_list:
            jet = _jet_source(residual_fn, None, x, policy.jet_order)
            certified.append(certified_local_error(jet, spacing))

    peak_location: float | None = None
    derivatives: list[float] | None = None
    if probe_jet is not None:
        peak_location, derivatives = float(probe_jet[0]), [float(v) for v in probe_jet[1]]
    elif policy.indicator in (Indicator.SINGULARITY, Indicator.SCALE_FLOW):
        values = residuals if certified is None else certified
        peak_i = max(range(len(values)), key=lambda i: abs(values[i]))
        peak_location = probe_list[peak_i]
        derivatives = _jet_source(residual_fn, target_fn, peak_location, policy.jet_order)

    working, report = refine_bank(
        packs,
        policy,
        probes=probe_list,
        residuals=residuals,
        contributions=contrib,
        derivatives=derivatives,
        certified_values=certified,
        peak_location=peak_location,
    )
    updated = bank_sync(bank, working)
    if before is not None and (report.born or report.grown) and not report.died:
        after = [float(v) for v in bank_forward(updated, probes_1d).reshape(-1)]
        assert_zero_perturbation(before, after)
    return updated, report

Related APIsomnibias.core.assert_zero_perturbation↗omnibias.core.refine.certified_local_error↗omnibias.core.refine.pack_contributions↗omnibias.core.refine_bank↗omnibias.jax.refine.bank_active_packs↗omnibias.jax.bank_forward↗omnibias.jax.refine.bank_sync↗omnibias.jax.refine.pack_term_norms↗

FUNCTION

scalar_jet_1d

omnibias.jax.refine.scalar_jet_1dSource ↗

scalar_jet_1d(fn: ScalarFn, x0: float, order: int) -> list[float]

phi^(k)(x0) for k = 0 .. order via nested jax.grad.

Parameters

Name	Type & default	Contract
fnpositional-or-keyword	ScalarFnRequired	See the operation contract and implementation below.
x0positional-or-keyword	floatRequired	See the operation contract and implementation below.
orderpositional-or-keyword	intRequired	See the operation contract and implementation below.

Returnslist[float]

Explicit validation & failure cases

• 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 · scalar jet exp

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_scalar_jet_exp() -> None:
    def fn(x: Array) -> Array:
        return jnp.exp(3.0 * x)

    derivs = scalar_jet_1d(fn, 0.0, 4)
    for k, val in enumerate(derivs):
        assert val == pytest.approx(3.0**k, rel=1e-10)

python -m pytest "packages/omnibias-jax/tests/test_refine.py::test_scalar_jet_exp" -q

Read the implementation

def scalar_jet_1d(fn: ScalarFn, x0: float, order: int) -> list[float]:
    """``phi^(k)(x0)`` for ``k = 0 .. order`` via nested ``jax.grad``."""
    if order < 0:
        raise ValueError(f"order must be >= 0, got {order}")

    def restriction(s: Array) -> Array:
        out = fn(jnp.reshape(s, ()))
        return jnp.reshape(out, ())

    x = jnp.asarray(float(x0))
    derivs = [float(restriction(x))]
    current: Callable[[Array], Array] = restriction
    for _ in range(order):
        current = jax.grad(current)
        derivs.append(float(current(x)))
    return derivs

← Package overviewRunnable learning examples →