API REFERENCE/omnibias-difference

omnibias.difference.jax

0.1.0a2AlphaPython >=3.10Apache-2.0

Finite-difference stencils, confluence limits, polynomial sequences, and scoped derivative-extraction certificates.

import omnibias.difference.jax

Module contract & mathematical conventions

JAX twin of the founding finite-difference stencil operator.

Bit-identical mirror of omnibias.difference.torch, built on the **shared pure-Python** stencil numbers in omnibias.difference._core.stencil. Pure and jit / vmap / grad-safe: name / order / delta / stencil are static, only z is traced, and no host-side coercion of a tracer occurs.

finite_difference_tower is the literal multi-bias sum sum_k s_k sigma(z + b_k) (numerical); collapse_to_derivative is the closed-form sigma^(order)(z) from the backend fastpath. Their residual -> 0 as delta -> 0 -- the founding bias collapse of many biases into the derivative sigma^(K-1).

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 →

FUNCTION

collapse_to_derivative

omnibias.difference.jax.collapse_to_derivativeSource ↗

collapse_to_derivative(name: str, z: Array, order: int) -> Array

Closed-form sigma^(order)(z) via the backend fastpath (the exact limit).

Parameters

Name	Type & default	Contract
namepositional-or-keyword	strRequired	See the operation contract and implementation below.
zpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
orderpositional-or-keyword	intRequired	See the operation contract and implementation below.

ReturnsArray

Explicit validation & failure cases

• When order < 0, raises ValueError(f'order must be >= 0, got {order}').
• When spec.fastpath is None, raises NotImplementedError(f'activation {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.

Executed regression · collapse to derivative parity

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_collapse_to_derivative_parity(name: str, order: int) -> None:
    zt = torch.tensor(_XS, dtype=torch.float64)
    zj = jnp.asarray(_XS)
    ct = DT.collapse_to_derivative(name, zt, order).detach().numpy()
    cj = np.asarray(DJ.collapse_to_derivative(name, zj, order))
    scale = 1.0 + float(np.max(np.abs(ct)))
    assert np.max(np.abs(ct - cj)) < 1e-8 * scale

python -m pytest "packages/omnibias-difference/tests/test_parity.py::test_collapse_to_derivative_parity[0-gaussian]" -q

Regression example · twins are jit and grad safe

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_twins_are_jit_and_grad_safe() -> None:
    zj = jnp.asarray(_XS)
    # jit through the numerical tower
    val = jax.jit(lambda z: DJ.finite_difference_tower("sigmoid", z, 1, 1e-2).sum())(zj)
    assert np.isfinite(float(val))
    # grad through the closed-form collapse
    grad = jax.grad(lambda z: DJ.collapse_to_derivative("tanh", z, 1).sum())(zj)
    assert grad.shape == zj.shape
    # torch autograd through the stencil operator
    zt = torch.tensor(_XS, dtype=torch.float64, requires_grad=True)
    DT.finite_difference_tower("sigmoid", zt, 1, 1e-2).sum().backward()
    assert zt.grad is not None and torch.isfinite(zt.grad).all()

python -m pytest packages/omnibias-difference/tests/test_parity.py -k test_twins_are_jit_and_grad_safe -q

Read the implementation

def collapse_to_derivative(name: str, z: Array, order: int) -> Array:
    """Closed-form ``sigma^(order)(z)`` via the backend fastpath (the exact limit)."""
    if order < 0:
        raise ValueError(f"order must be >= 0, got {order}")
    spec = get_activation(name)
    if spec.fastpath is None:
        raise NotImplementedError(f"activation {name!r} has no closed-form derivative kernel")
    out: Array = spec.fastpath(z, order)
    return out

Related APIsomnibias.jax.get_activation↗

FUNCTION

finite_difference_residual

omnibias.difference.jax.finite_difference_residualSource ↗

finite_difference_residual(name: str, z: Array, order: int, delta: float, stencil: Stencil='central') -> Array

finite_difference_tower - collapse_to_derivative; -> 0 as delta -> 0.

Parameters

Name	Type & default	Contract
namepositional-or-keyword	strRequired	See the operation contract and implementation below.
zpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
orderpositional-or-keyword	intRequired	See the operation contract and implementation below.
deltapositional-or-keyword	floatRequired	See the operation contract and implementation below.
stencilpositional-or-keyword	StencilDefault: 'central'	See the operation contract and implementation below.

ReturnsArray

Executed regression · residual parity

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_residual_parity(name: str) -> None:
    delta = 1e-2
    zt = torch.tensor(_XS, dtype=torch.float64)
    zj = jnp.asarray(_XS)
    rt = DT.finite_difference_residual(name, zt, 2, delta, "central").detach().numpy()
    rj = np.asarray(DJ.finite_difference_residual(name, zj, 2, delta, "central"))
    sign_mag = sum(abs(s) for s in stencil_signs(2, delta, "central"))
    assert np.max(np.abs(rt - rj)) < 1e-12 * sign_mag + 1e-12

python -m pytest "packages/omnibias-difference/tests/test_parity.py::test_residual_parity[gaussian]" -q

Read the implementation

def finite_difference_residual(
    name: str, z: Array, order: int, delta: float, stencil: Stencil = "central"
) -> Array:
    """``finite_difference_tower - collapse_to_derivative``; ``-> 0`` as ``delta -> 0``."""
    return finite_difference_tower(name, z, order, delta, stencil) - collapse_to_derivative(
        name, z, order
    )

Related APIsomnibias.difference.jax.collapse_to_derivative↗omnibias.difference.jax.finite_difference_tower↗

FUNCTION

finite_difference_tower

omnibias.difference.jax.finite_difference_towerSource ↗

finite_difference_tower(name: str, z: Array, order: int, delta: float, stencil: Stencil='central') -> Array

Evaluate the numerical multi-bias stencil using the JAX backend.

The shared stencil defines offsets, signs, and normalization. This is finite-gap evaluation rather than the analytic zero-gap derivative. Small offsets can amplify floating-point cancellation. Shared coefficients do not guarantee bitwise-equal Torch and JAX results; validate dtype, the selected gap, derivative order, and activation path.

Parameters

Name	Type & default	Contract
namepositional-or-keyword	strRequired	See the operation contract and implementation below.
zpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
orderpositional-or-keyword	intRequired	See the operation contract and implementation below.
deltapositional-or-keyword	floatRequired	See the operation contract and implementation below.
stencilpositional-or-keyword	StencilDefault: 'central'	See the operation contract and implementation below.

ReturnsArray

Executed regression · residual parity

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_residual_parity(name: str) -> None:
    delta = 1e-2
    zt = torch.tensor(_XS, dtype=torch.float64)
    zj = jnp.asarray(_XS)
    rt = DT.finite_difference_residual(name, zt, 2, delta, "central").detach().numpy()
    rj = np.asarray(DJ.finite_difference_residual(name, zj, 2, delta, "central"))
    sign_mag = sum(abs(s) for s in stencil_signs(2, delta, "central"))
    assert np.max(np.abs(rt - rj)) < 1e-12 * sign_mag + 1e-12

python -m pytest "packages/omnibias-difference/tests/test_parity.py::test_residual_parity[gaussian]" -q

Regression example · finite difference tower 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_finite_difference_tower_parity(name: str, order: int, stencil: str) -> None:
    delta = 1e-2
    zt = torch.tensor(_XS, dtype=torch.float64)
    zj = jnp.asarray(_XS)
    ft = DT.finite_difference_tower(name, zt, order, delta, stencil).detach().numpy()
    fj = np.asarray(DJ.finite_difference_tower(name, zj, order, delta, stencil))
    sign_mag = sum(abs(s) for s in stencil_signs(order, delta, stencil))
    tol = 1e-12 * sign_mag + 1e-12  # calibrated to the 1/delta^order amplification
    assert np.max(np.abs(ft - fj)) < tol

python -m pytest packages/omnibias-difference/tests/test_parity.py -k test_finite_difference_tower_parity -q

Read the implementation

def finite_difference_tower(
    name: str, z: Array, order: int, delta: float, stencil: Stencil = "central"
) -> Array:
    """Numerical multi-bias finite-difference stencil ``sum_k s_k sigma(z + b_k)``.

    The signs and offsets come from the shared pure-Python
    :mod:`omnibias.difference._core.stencil`, so the value is bit-identical to the
    PyTorch twin up to the base activation's per-backend rounding.
    """
    spec = get_activation(name)
    signs = jnp.asarray(stencil_signs(order, delta, stencil), dtype=z.dtype)
    offsets = jnp.asarray(stencil_offsets(order, delta, stencil), dtype=z.dtype)
    values = spec.forward(z[..., None] + offsets)  # (..., K)
    return (values * signs).sum(axis=-1)

Related APIsomnibias.difference.stencil_offsets↗omnibias.difference.stencil_signs↗omnibias.jax.get_activation↗

← Package overviewRunnable learning examples →