omnibias.jax.multipack

0.5.0rc1Release candidatePython >=3.10Apache-2.0

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

import omnibias.jax.multipack

Module contract & mathematical conventions

Heterogeneous multi-pack Birkhoff unit (JAX twin of torch; theory 01-01).

Closed-form evaluation of

F(z) = sum_g c_g * sigma^(n_g)(z + mu_g)

Orders and mean_index are static Python tuples so the path stays jit-safe. Founding bias collapse delta -> 0 only -- no temperature collapse.

Shared formulas do not imply universal floating-point identity. Verify dtype, shape, derivative-order support, and the stated numerical assumptions for the path you use.

Module source ↗Source revision 9a2d73df0573Running reference examples →

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
`MultiPackSpec`	`omnibias.core.MultiPackSpec`
`PackSpec`	`omnibias.core.PackSpec`

FUNCTION

init_multipack

omnibias.jax.init_multipackSource ↗

init_multipack(num_channels: int, spec: MultiPackSpec, *, base: str | JaxActivationSpec='sigmoid', share_means: bool=True) -> tuple[JaxActivationSpec, Array, Array, tuple[int, ...], tuple[int, ...]]

Return (act, means, weights, orders, mean_index).

Parameters

Name	Type & default	Contract
`num_channels`positional-or-keyword	`int`Required	See the operation contract and implementation below.
`spec`positional-or-keyword	`MultiPackSpec`Required	See the operation contract and implementation below.
`base`keyword-only	`str \| JaxActivationSpec`Default: 'sigmoid'	See the operation contract and implementation below.
`share_means`keyword-only	`bool`Default: True	See the operation contract and implementation below.

Returnstuple[JaxActivationSpec, Array, Array, tuple[int, ...], tuple[int, ...]]

Explicit validation & failure cases

When num_channels < 1, raises ValueError(f'num_channels must be >= 1, got {num_channels}').
When act.fastpath is None, raises NotImplementedError(f'Activation {act.name!r} has no closed-form derivative kernel').

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

Regression example · torch jax multipack 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_torch_jax_multipack_parity(shared_activation: bool, monkeypatch) -> None:
    import jax
    import jax.numpy as jnp
    from omnibias.jax.multipack import init_multipack, multipack_apply
    from omnibias.torch.multipack import MultiPackUnit

    jax.config.update("jax_enable_x64", True)
    torch.set_default_dtype(torch.float64)
    if shared_activation:
        # Isolate shared derivative-tower arithmetic from platform libm differences.
        def shared_tanh(values):
            array = np.asarray(values)
            with mpmath.workdps(80):
                return np.array([
                    float(mpmath.tanh(mpmath.mpf(float(x)))) for x in array.ravel()
                ]).reshape(array.shape)

        monkeypatch.setattr(
            torch, "tanh",
            lambda values: torch.as_tensor(
                shared_tanh(values.detach().numpy()), dtype=values.dtype, device=values.device
            ),
        )
        monkeypatch.setattr(jnp, "tanh", lambda values: jnp.asarray(shared_tanh(values)))
    spec = MultiPackSpec((PackSpec(1, -0.5, 1.0), PackSpec(2, 0.5, 0.25)))
    unit = MultiPackUnit(
        2, spec, base="tanh", learnable_means=False, learnable_weights=False
    )
    z_np = np.array([[0.0, 0.0], [0.25, -0.25], [0.5, 0.5]], dtype=np.float64)
    torch_out = unit(torch.as_tensor(z_np)).detach().numpy()
    _act, means, weights, orders, mean_index = init_multipack(
        2, spec, base="tanh", share_means=True
    )
    jax_out = np.asarray(
        multipack_apply(
            jnp.asarray(z_np), means, weights, orders, "tanh", mean_index=mean_index
        )
    )
    with mpmath.workdps(80):
        reference = np.array([
            [float(sum(
                mpmath.mpf(p.weight) * mpmath.diff(
                    mpmath.tanh, mpmath.mpf(float(z)) + mpmath.mpf(p.mean), p.order
                )
                for p in spec.packs
            )) for z in row]
            for row in z_np
        ])
    # The native tanh kernels can differ by a ULP even away from their tails.
    # Both outputs must independently retain the established four-ULP accuracy.
    np.testing.assert_array_max_ulp(torch_out, reference, maxulp=4)
    np.testing.assert_array_max_ulp(jax_out, reference, maxulp=4)
    if shared_activation:
        np.testing.assert_array_equal(torch_out, jax_out)

python -m pytest tests/test_multipack_parity.py -k test_torch_jax_multipack_parity -q

Read the implementation

def init_multipack(
    num_channels: int,
    spec: MultiPackSpec,
    *,
    base: str | JaxActivationSpec = "sigmoid",
    share_means: bool = True,
) -> tuple[JaxActivationSpec, Array, Array, tuple[int, ...], tuple[int, ...]]:
    """Return ``(act, means, weights, orders, mean_index)``."""
    if num_channels < 1:
        raise ValueError(f"num_channels must be >= 1, got {num_channels}")
    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"
        )
    orders = tuple(int(p.order) for p in spec.packs)
    probe = jnp.zeros(())
    for n in orders:
        act.fastpath(probe, n)
    if share_means:
        distinct = spec.distinct_means
        mean_to_slot = {m: i for i, m in enumerate(distinct)}
        mean_index = tuple(mean_to_slot[p.mean] for p in spec.packs)
        means = jnp.asarray(list(distinct), dtype=jnp.float64)
    else:
        mean_index = tuple(range(len(spec.packs)))
        means = jnp.asarray([p.mean for p in spec.packs], dtype=jnp.float64)
    weights = jnp.asarray([p.weight for p in spec.packs], dtype=jnp.float64)
    return act, means, weights, orders, mean_index

Other public import paths

omnibias.jax.init_multipack
omnibias.jax.multipack.init_multipack

FUNCTION

multipack_apply

omnibias.jax.BirkhoffOMBUSource ↗

multipack_apply(z: Array, means: Array, weights: Array, orders: tuple[int, ...], base: str | JaxActivationSpec='sigmoid', *, mean_index: tuple[int, ...] | None=None) -> Array

Functional multi-pack forward (mathematically equivalent twin of MultiPackUnit).

Backend scope: this identifies the corresponding mathematical construction. It does not guarantee bitwise-equal floating-point results across devices, compilers, dtypes, or backend versions. Validate the selected path with an appropriate absolute and relative tolerance.

Parameters

Name	Type & default	Contract
`z`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`means`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`weights`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`orders`positional-or-keyword	`tuple[int, ...]`Required	See the operation contract and implementation below.
`base`positional-or-keyword	`str \| JaxActivationSpec`Default: 'sigmoid'	See the operation contract and implementation below.
`mean_index`keyword-only	`tuple[int, ...] \| None`Default: None	See the operation contract and implementation below.

ReturnsArray

Regression example · torch jax multipack 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_torch_jax_multipack_parity(shared_activation: bool, monkeypatch) -> None:
    import jax
    import jax.numpy as jnp
    from omnibias.jax.multipack import init_multipack, multipack_apply
    from omnibias.torch.multipack import MultiPackUnit

    jax.config.update("jax_enable_x64", True)
    torch.set_default_dtype(torch.float64)
    if shared_activation:
        # Isolate shared derivative-tower arithmetic from platform libm differences.
        def shared_tanh(values):
            array = np.asarray(values)
            with mpmath.workdps(80):
                return np.array([
                    float(mpmath.tanh(mpmath.mpf(float(x)))) for x in array.ravel()
                ]).reshape(array.shape)

        monkeypatch.setattr(
            torch, "tanh",
            lambda values: torch.as_tensor(
                shared_tanh(values.detach().numpy()), dtype=values.dtype, device=values.device
            ),
        )
        monkeypatch.setattr(jnp, "tanh", lambda values: jnp.asarray(shared_tanh(values)))
    spec = MultiPackSpec((PackSpec(1, -0.5, 1.0), PackSpec(2, 0.5, 0.25)))
    unit = MultiPackUnit(
        2, spec, base="tanh", learnable_means=False, learnable_weights=False
    )
    z_np = np.array([[0.0, 0.0], [0.25, -0.25], [0.5, 0.5]], dtype=np.float64)
    torch_out = unit(torch.as_tensor(z_np)).detach().numpy()
    _act, means, weights, orders, mean_index = init_multipack(
        2, spec, base="tanh", share_means=True
    )
    jax_out = np.asarray(
        multipack_apply(
            jnp.asarray(z_np), means, weights, orders, "tanh", mean_index=mean_index
        )
    )
    with mpmath.workdps(80):
        reference = np.array([
            [float(sum(
                mpmath.mpf(p.weight) * mpmath.diff(
                    mpmath.tanh, mpmath.mpf(float(z)) + mpmath.mpf(p.mean), p.order
                )
                for p in spec.packs
            )) for z in row]
            for row in z_np
        ])
    # The native tanh kernels can differ by a ULP even away from their tails.
    # Both outputs must independently retain the established four-ULP accuracy.
    np.testing.assert_array_max_ulp(torch_out, reference, maxulp=4)
    np.testing.assert_array_max_ulp(jax_out, reference, maxulp=4)
    if shared_activation:
        np.testing.assert_array_equal(torch_out, jax_out)

python -m pytest tests/test_multipack_parity.py -k test_torch_jax_multipack_parity -q

Read the implementation

def multipack_apply(
    z: Array,
    means: Array,
    weights: Array,
    orders: tuple[int, ...],
    base: str | JaxActivationSpec = "sigmoid",
    *,
    mean_index: tuple[int, ...] | None = None,
) -> Array:
    """Functional multi-pack forward (bit-identical twin of ``MultiPackUnit``)."""
    act = get_activation(base) if isinstance(base, str) else base
    return multipack_response(
        z, means, weights, orders, act, mean_index=mean_index
    )

Other public import paths

omnibias.jax.BirkhoffOMBU
omnibias.jax.multipack_apply
omnibias.jax.multipack.BirkhoffOMBU
omnibias.jax.multipack.multipack_apply

FUNCTION

multipack_response

omnibias.jax.multipack_responseSource ↗

multipack_response(z: Array, means: Array, weights: Array, orders: tuple[int, ...], spec: JaxActivationSpec, *, mean_index: tuple[int, ...] | None=None) -> Array

Evaluate sum_g c_g sigma^(n_g)(z + mu_g) with a closed-form tower.

orders / mean_index must be Python tuples of static ints.

Parameters

Name	Type & default	Contract
`z`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`means`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`weights`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`orders`positional-or-keyword	`tuple[int, ...]`Required	See the operation contract and implementation below.
`spec`positional-or-keyword	`JaxActivationSpec`Required	See the operation contract and implementation below.
`mean_index`keyword-only	`tuple[int, ...] \| None`Default: None	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 int(weights_a.shape[-1]) != G, raises ValueError('weights length must equal number of packs').
When int(means_a.shape[-1]) != G, raises ValueError('means length must equal number of packs').
When len(mean_index) != G, raises ValueError('mean_index length must equal number of packs').
When n < 0, raises ValueError(f'order must be >= 0, got {n}').

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

Read the implementation

def multipack_response(
    z: Array,
    means: Array,
    weights: Array,
    orders: tuple[int, ...],
    spec: JaxActivationSpec,
    *,
    mean_index: tuple[int, ...] | None = None,
) -> Array:
    """Evaluate ``sum_g c_g sigma^(n_g)(z + mu_g)`` with a closed-form tower.

    ``orders`` / ``mean_index`` must be Python tuples of static ints.
    """
    if spec.fastpath is None:
        raise NotImplementedError(
            f"Activation {spec.name!r} has no closed-form derivative kernel"
        )
    means_a = jnp.asarray(means)
    weights_a = jnp.asarray(weights)
    G = len(orders)
    if int(weights_a.shape[-1]) != G:
        raise ValueError("weights length must equal number of packs")
    if mean_index is None:
        if int(means_a.shape[-1]) != G:
            raise ValueError("means length must equal number of packs")
        index = tuple(range(G))
    else:
        if len(mean_index) != G:
            raise ValueError("mean_index length must equal number of packs")
        index = mean_index
    for n in orders:
        if n < 0:
            raise ValueError(f"order must be >= 0, got {n}")

    fp = spec.fastpath
    slots = sorted(set(index))
    slot_u = {s: z + means_a[s] for s in slots}
    towers: dict[int, Array] = {}
    if spec.tower is not None:
        for slot in slots:
            highest = max(n for g, n in enumerate(orders) if index[g] == slot)
            towers[slot] = spec.tower(slot_u[slot], highest)
    out: Array | None = None
    for g, n in enumerate(orders):
        value = towers[index[g]][n] if spec.tower is not None else fp(slot_u[index[g]], n)
        term = weights_a[g] * value
        out = term if out is None else out + term
    assert out is not None
    return out

Other public import paths

omnibias.jax.multipack_response
omnibias.jax.multipack.multipack_response

FUNCTION

packs_from_arrays

omnibias.jax.multipack.packs_from_arraysSource ↗

packs_from_arrays(orders: Sequence[int], means: Sequence[float], weights: Sequence[float] | None=None) -> MultiPackSpec

Build a MultiPackSpec from parallel sequences.

Parameters

Name	Type & default	Contract
`orders`positional-or-keyword	`Sequence[int]`Required	See the operation contract and implementation below.
`means`positional-or-keyword	`Sequence[float]`Required	See the operation contract and implementation below.
`weights`positional-or-keyword	`Sequence[float] \| None`Default: None	See the operation contract and implementation below.

ReturnsMultiPackSpec

Explicit validation & failure cases

When not len(orders) == len(means) == len(w), raises ValueError('orders, means, weights length mismatch').

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

Read the implementation

def packs_from_arrays(
    orders: Sequence[int],
    means: Sequence[float],
    weights: Sequence[float] | None = None,
) -> MultiPackSpec:
    """Build a :class:`MultiPackSpec` from parallel sequences."""
    w = (1.0,) * len(orders) if weights is None else tuple(float(x) for x in weights)
    if not (len(orders) == len(means) == len(w)):
        raise ValueError("orders, means, weights length mismatch")
    return MultiPackSpec(
        tuple(
            PackSpec(order=int(n), mean=float(mu), weight=float(c))
            for n, mu, c in zip(orders, means, w, strict=True)
        )
    )

← Package overview Runnable learning examples →

omnibias.jax.multipack

Imports & aliases​

init_multipack​

Parameters

multipack_apply​

Parameters

multipack_response​

Parameters

packs_from_arrays​

Parameters

Imports & aliases

init_multipack

multipack_apply

multipack_response

packs_from_arrays