omnibias.jax.scan

0.5.0rc1Release candidatePython >=3.10Apache-2.0

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

import omnibias.jax.scan

Module contract & mathematical conventions

Bias scan: shared template on a bank of offsets (JAX twin; theory 01-02).

Functional twin of omnibias.torch.scan. Orders stay Python tuples so the path is jit-safe. Equivariance is an interior lattice shift, not a circular wrap. Soft-argmax gamma is not bias 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
`BankSpec`	`omnibias.core.BankSpec`

FUNCTION

bias_scan

omnibias.jax.bias_scanSource ↗

bias_scan(z: Array, offsets: Array, scales: Array, template: MultiPackSpec | str, base: str | JaxActivationSpec='tanh', *, gap: float=1.0, readout: Readout='response', gamma: float=8.0, pool_taps: Array | None=None) -> Array

Functional scan forward (mathematically equivalent twin of omnibias.torch.scan.BiasScan).

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.
`offsets`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`scales`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`template`positional-or-keyword	`MultiPackSpec \| str`Required	See the operation contract and implementation below.
`base`positional-or-keyword	`str \| JaxActivationSpec`Default: 'tanh'	See the operation contract and implementation below.
`gap`keyword-only	`float`Default: 1.0	See the operation contract and implementation below.
`readout`keyword-only	`Readout`Default: 'response'	See the operation contract and implementation below.
`gamma`keyword-only	`float`Default: 8.0	See the operation contract and implementation below.
`pool_taps`keyword-only	`Array \| None`Default: None	See the operation contract and implementation below.

ReturnsArray

Explicit validation & failure cases

When readout == 'pooled', raises ValueError('pooled readout shape mismatch').

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

Regression example · g3 argmax within 4 ulp

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_argmax_within_4_ulp() -> None:
    import jax
    import jax.numpy as jnp
    from omnibias.jax.scan import bias_scan, init_bias_scan
    from omnibias.torch.scan import BiasScan

    jax.config.update("jax_enable_x64", True)
    torch.set_default_dtype(torch.float64)
    bank = BankSpec((-1.0, -0.5, 0.0, 0.5, 1.0))
    unit = BiasScan(
        1,
        bank,
        template="grad",
        base="tanh",
        learnable_offsets=False,
        learnable_scales=False,
        readout="argmax",
        gamma=8.0,
        dtype=torch.float64,
    )
    z_np = np.array([[0.3], [0.0], [-0.2]], dtype=np.float64)
    torch_out = unit(torch.as_tensor(z_np)).detach().numpy().reshape(-1)
    act, offsets, scales, tmpl, _taps = init_bias_scan(1, bank, template="grad", base="tanh")
    jax_out = np.asarray(
        bias_scan(
            jnp.asarray(z_np),
            offsets,
            scales,
            tmpl,
            act,
            readout="argmax",
            gamma=8.0,
        )
    ).reshape(-1)
    worst = max(_ulp_error(float(a), float(b)) for a, b in zip(torch_out, jax_out, strict=True))
    assert worst <= 4.0, f"argmax parity worst_ulp={worst}"

python -m pytest tests/test_scan_parity.py -k test_g3_argmax_within_4_ulp -q

Regression example · g3 torch jax bit identical response

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_torch_jax_bit_identical_response() -> None:
    import jax
    import jax.numpy as jnp
    from omnibias.jax.scan import bias_scan, init_bias_scan
    from omnibias.torch.scan import BiasScan

    jax.config.update("jax_enable_x64", True)
    torch.set_default_dtype(torch.float64)
    bank = BankSpec.uniform(-1.0, 1.0, 5)
    unit = BiasScan(
        2,
        bank,
        template="grad",
        base="tanh",
        learnable_offsets=False,
        learnable_scales=False,
        readout="response",
        dtype=torch.float64,
    )
    z_np = np.array([[0.0, 0.0], [0.3, -0.25], [0.5, 0.5]], dtype=np.float64)
    torch_out = unit(torch.as_tensor(z_np)).detach().numpy()
    act, offsets, scales, tmpl, _taps = init_bias_scan(2, bank, template="grad", base="tanh")
    jax_out = np.asarray(
        bias_scan(
            jnp.asarray(z_np),
            offsets,
            scales,
            tmpl,
            act,
            readout="response",
        )
    )
    worst = max(
        _ulp_error(float(a), float(b))
        for a, b in zip(torch_out.reshape(-1), jax_out.reshape(-1), strict=True)
    )
    assert worst <= 4.0, f"response parity worst_ulp={worst}"

python -m pytest tests/test_scan_parity.py -k test_g3_torch_jax_bit_identical_response -q

Read the implementation

def bias_scan(
    z: Array,
    offsets: Array,
    scales: Array,
    template: MultiPackSpec | str,
    base: str | JaxActivationSpec = "tanh",
    *,
    gap: float = 1.0,
    readout: Readout = "response",
    gamma: float = 8.0,
    pool_taps: Array | None = None,
) -> Array:
    """Functional scan forward (bit-identical twin of :class:`omnibias.torch.scan.BiasScan`)."""
    act = get_activation(base) if isinstance(base, str) else base
    resp = scan_response(z, offsets, scales, template, act, gap=gap)
    if readout == "response":
        return resp
    taps = (
        jnp.full((offsets.size,), 1.0 / max(int(offsets.size), 1), dtype=resp.dtype)
        if pool_taps is None
        else pool_taps
    )
    if readout == "pooled":
        if resp.shape[-1] == taps.size:
            return (resp * taps).sum(axis=-1)
        if resp.ndim >= 2 and resp.shape[-2] == taps.size:
            return (resp * taps.reshape((-1, 1))).sum(axis=-2).mean(axis=-1)
        raise ValueError("pooled readout shape mismatch")
    loc = soft_argmax_offset(resp, offsets, gamma=gamma)
    if loc.ndim == z.ndim:
        return loc
    return loc.mean(axis=-1)

Other public import paths

omnibias.jax.bias_scan
omnibias.jax.scan.bias_scan

FUNCTION

init_bias_scan

omnibias.jax.init_bias_scanSource ↗

init_bias_scan(num_channels: int, bank: BankSpec, *, template: MultiPackSpec | OpName | str | None=None, op: OpName | str | None=None, base: str | JaxActivationSpec='tanh', derivative_order: int=1, gap: float=1.0) -> tuple[JaxActivationSpec, Array, Array, MultiPackSpec | str, Array]

Return (act, offsets, scales, template, pool_taps).

op is the 01-13 catalog alias for template. Pass only one.

Parameters

Name	Type & default	Contract
`num_channels`positional-or-keyword	`int`Required	See the operation contract and implementation below.
`bank`positional-or-keyword	`BankSpec`Required	See the operation contract and implementation below.
`template`keyword-only	`MultiPackSpec \| OpName \| str \| None`Default: None	See the operation contract and implementation below.
`op`keyword-only	`OpName \| str \| None`Default: None	See the operation contract and implementation below.
`base`keyword-only	`str \| JaxActivationSpec`Default: 'tanh'	See the operation contract and implementation below.
`derivative_order`keyword-only	`int`Default: 1	See the operation contract and implementation below.
`gap`keyword-only	`float`Default: 1.0	See the operation contract and implementation below.

Returnstuple[JaxActivationSpec, Array, Array, MultiPackSpec | str, Array]

Explicit validation & failure cases

When num_channels < 1, raises ValueError(f'num_channels must be >= 1, got {num_channels}').
When op is not None, raises ValueError('pass only one of template= or op=; they name the same role').

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

Regression example · g3 argmax within 4 ulp

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_argmax_within_4_ulp() -> None:
    import jax
    import jax.numpy as jnp
    from omnibias.jax.scan import bias_scan, init_bias_scan
    from omnibias.torch.scan import BiasScan

    jax.config.update("jax_enable_x64", True)
    torch.set_default_dtype(torch.float64)
    bank = BankSpec((-1.0, -0.5, 0.0, 0.5, 1.0))
    unit = BiasScan(
        1,
        bank,
        template="grad",
        base="tanh",
        learnable_offsets=False,
        learnable_scales=False,
        readout="argmax",
        gamma=8.0,
        dtype=torch.float64,
    )
    z_np = np.array([[0.3], [0.0], [-0.2]], dtype=np.float64)
    torch_out = unit(torch.as_tensor(z_np)).detach().numpy().reshape(-1)
    act, offsets, scales, tmpl, _taps = init_bias_scan(1, bank, template="grad", base="tanh")
    jax_out = np.asarray(
        bias_scan(
            jnp.asarray(z_np),
            offsets,
            scales,
            tmpl,
            act,
            readout="argmax",
            gamma=8.0,
        )
    ).reshape(-1)
    worst = max(_ulp_error(float(a), float(b)) for a, b in zip(torch_out, jax_out, strict=True))
    assert worst <= 4.0, f"argmax parity worst_ulp={worst}"

python -m pytest tests/test_scan_parity.py -k test_g3_argmax_within_4_ulp -q

Regression example · g3 torch jax bit identical response

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_torch_jax_bit_identical_response() -> None:
    import jax
    import jax.numpy as jnp
    from omnibias.jax.scan import bias_scan, init_bias_scan
    from omnibias.torch.scan import BiasScan

    jax.config.update("jax_enable_x64", True)
    torch.set_default_dtype(torch.float64)
    bank = BankSpec.uniform(-1.0, 1.0, 5)
    unit = BiasScan(
        2,
        bank,
        template="grad",
        base="tanh",
        learnable_offsets=False,
        learnable_scales=False,
        readout="response",
        dtype=torch.float64,
    )
    z_np = np.array([[0.0, 0.0], [0.3, -0.25], [0.5, 0.5]], dtype=np.float64)
    torch_out = unit(torch.as_tensor(z_np)).detach().numpy()
    act, offsets, scales, tmpl, _taps = init_bias_scan(2, bank, template="grad", base="tanh")
    jax_out = np.asarray(
        bias_scan(
            jnp.asarray(z_np),
            offsets,
            scales,
            tmpl,
            act,
            readout="response",
        )
    )
    worst = max(
        _ulp_error(float(a), float(b))
        for a, b in zip(torch_out.reshape(-1), jax_out.reshape(-1), strict=True)
    )
    assert worst <= 4.0, f"response parity worst_ulp={worst}"

python -m pytest tests/test_scan_parity.py -k test_g3_torch_jax_bit_identical_response -q

Read the implementation

def init_bias_scan(
    num_channels: int,
    bank: BankSpec,
    *,
    template: MultiPackSpec | OpName | str | None = None,
    op: OpName | str | None = None,
    base: str | JaxActivationSpec = "tanh",
    derivative_order: int = 1,
    gap: float = 1.0,
) -> tuple[JaxActivationSpec, Array, Array, MultiPackSpec | str, Array]:
    """Return ``(act, offsets, scales, template, pool_taps)``.

    ``op`` is the 01-13 catalog alias for ``template``. Pass only one.
    """
    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
    tmpl: MultiPackSpec | str
    if isinstance(template, MultiPackSpec):
        if op is not None:
            raise ValueError("pass only one of template= or op=; they name the same role")
        tmpl = template
    else:
        tmpl = template_from_op(
            resolve_scan_role(template, op),
            derivative_order=derivative_order,
            gap=gap,
        )
    offsets = jnp.asarray(list(bank.offsets), dtype=jnp.float64)
    scales = jnp.asarray(list(bank.scales), dtype=jnp.float64)
    n = int(offsets.size)
    pool_taps = jnp.full((n,), 1.0 / max(n, 1), dtype=jnp.float64)
    return act, offsets, scales, tmpl, pool_taps

Other public import paths

omnibias.jax.init_bias_scan
omnibias.jax.scan.init_bias_scan

FUNCTION

min_offset_separation

omnibias.jax.scan.min_offset_separationSource ↗

min_offset_separation(offsets: Array) -> Array

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
`offsets`positional-or-keyword	`Array`Required	See the operation contract and implementation below.

ReturnsArray

Read the implementation

def min_offset_separation(offsets: Array) -> Array:
    ordered = jnp.sort(jnp.reshape(offsets, (-1,)))
    if ordered.size < 2:
        return jnp.asarray(float("inf"), dtype=offsets.dtype)
    return jnp.min(ordered[1:] - ordered[:-1])

FUNCTION

scan_response

omnibias.jax.scan_responseSource ↗

scan_response(z: Array, offsets: Array, scales: Array, spec: MultiPackSpec | str, base: JaxActivationSpec, *, gap: float=1.0) -> Array

Evaluate the template at scale * (z + offset).

Parameters

Name	Type & default	Contract
`z`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`offsets`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`scales`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`spec`positional-or-keyword	`MultiPackSpec \| str`Required	See the operation contract and implementation below.
`base`positional-or-keyword	`JaxActivationSpec`Required	See the operation contract and implementation below.
`gap`keyword-only	`float`Default: 1.0	See the operation contract and implementation below.

ReturnsArray

Explicit validation & failure cases

When base.forward is None, raises NotImplementedError('band template needs ActivationSpec.forward').
When base.integral is None, raises NotImplementedError(f'Activation {base.name!r} has no closed-form integral kernel').

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

Read the implementation

def scan_response(
    z: Array,
    offsets: Array,
    scales: Array,
    spec: MultiPackSpec | str,
    base: JaxActivationSpec,
    *,
    gap: float = 1.0,
) -> Array:
    """Evaluate the template at ``scale * (z + offset)``."""
    off = jnp.reshape(offsets, (-1,))
    sc = jnp.reshape(scales, (-1,))
    bank = z[..., None] + off
    extra_scale = int(sc.size) != 1
    if extra_scale:
        u = bank[..., None] * sc
    else:
        u = sc.reshape(()) * bank
    if isinstance(spec, str):
        half = 0.5 * float(gap)
        if spec == "band":
            if base.forward is None:
                raise NotImplementedError("band template needs ActivationSpec.forward")
            out = base.forward(u + half) - base.forward(u - half)
        elif spec == "integral":
            if base.integral is None:
                raise NotImplementedError(
                    f"Activation {base.name!r} has no closed-form integral kernel"
                )
            out = base.integral(u + half) - base.integral(u - half)
        else:
            raise ValueError(f"unknown window template {spec!r}")
    else:
        orders = tuple(int(p.order) for p in spec.packs)
        means = jnp.asarray([p.mean for p in spec.packs], dtype=u.dtype)
        weights = jnp.asarray([p.weight for p in spec.packs], dtype=u.dtype)
        out = multipack_response(u, means, weights, orders, base)
    return out

Other public import paths

omnibias.jax.scan_response
omnibias.jax.scan.scan_response

FUNCTION

soft_argmax_offset

omnibias.jax.soft_argmax_offsetSource ↗

soft_argmax_offset(response: Array, offsets: Array, *, gamma: float=8.0) -> Array

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
`response`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`offsets`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`gamma`keyword-only	`float`Default: 8.0	See the operation contract and implementation below.

ReturnsArray

Read the implementation

def soft_argmax_offset(response: Array, offsets: Array, *, gamma: float = 8.0) -> Array:
    off = jnp.reshape(offsets, (-1,))
    if response.shape[-1] == off.size:
        w = softmax(float(gamma) * response, axis=-1)
        return (w * off).sum(axis=-1)
    if response.ndim >= 2 and response.shape[-2] == off.size:
        w = softmax(float(gamma) * response, axis=-2)
        return (w * off.reshape((-1, 1))).sum(axis=-2)
    raise ValueError("response trailing shape does not match offsets")

Other public import paths

omnibias.jax.soft_argmax_offset
omnibias.jax.scan.soft_argmax_offset

FUNCTION

template_from_op

omnibias.jax.scan.template_from_opSource ↗

template_from_op(op: OpName | str, *, derivative_order: int=1, gap: float=1.0) -> MultiPackSpec | str

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
`op`positional-or-keyword	`OpName \| str`Required	See the operation contract and implementation below.
`derivative_order`keyword-only	`int`Default: 1	See the operation contract and implementation below.
`gap`keyword-only	`float`Default: 1.0	See the operation contract and implementation below.

ReturnsMultiPackSpec | str

Explicit validation & failure cases

When int(derivative_order) < 0, raises ValueError('derivative_order must be >= 0').
When gap <= 0.0, raises ValueError('gap must be positive').

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

Read the implementation

def template_from_op(
    op: OpName | str,
    *,
    derivative_order: int = 1,
    gap: float = 1.0,
) -> MultiPackSpec | str:
    name = str(op).lower()
    if name in _COLLAPSE_ORDER:
        return MultiPackSpec((PackSpec(_COLLAPSE_ORDER[name], 0.0),))
    if name == "derivative":
        if int(derivative_order) < 0:
            raise ValueError("derivative_order must be >= 0")
        return MultiPackSpec((PackSpec(int(derivative_order), 0.0),))
    if name in ("band", "integral"):
        if gap <= 0.0:
            raise ValueError("gap must be positive")
        return name
    raise ValueError(f"unknown template op {op!r}")

← Package overview Runnable learning examples →

omnibias.jax.scan

Imports & aliases​

bias_scan​

Parameters

init_bias_scan​

Parameters

min_offset_separation​

Parameters

scan_response​

Parameters

soft_argmax_offset​

Parameters

template_from_op​

Parameters

Imports & aliases

bias_scan

init_bias_scan

min_offset_separation

scan_response

soft_argmax_offset

template_from_op