Skip to main content

omnibias.torch.fastpath.eulerian

0.5.0rc1Release candidatePython >=3.10Apache-2.0

PyTorch activation derivatives, trainable operator layers, network jets, and direct differential-operator contractions.

import omnibias.torch.fastpath.eulerian
Module contract & mathematical conventions

Closed-form derivatives of sigmoid (and softplus) via the Eulerian/sigmoid-polynomial recursion.

Let s = sigma(z) = 1 / (1 + e^{-z}). Then s' = s (1 - s) (a Riccati identity), and every higher derivative of sigma is a polynomial in s:

sigma^(n)(z) = P_n(s), where P_0(s) = s P_{n+1}(s) = s (1 - s) * P_n'(s).

The coefficients of P_n are (signed) Eulerian numbers: P_n(s) = sum_{k=0}^{n} (-1)^{n-k} <n, k> s^{k+1} (1 - s)^{n-k}. The recurrence lives in the **canonical** omnibias.core.polynomials module so the torch and JAX backends always share a single source of truth. We import it here and re-export it under the historical name for back-compat (callers writing from omnibias.torch.fastpath.eulerian import sigmoid_polynomial_coeffs keep working).

For softplus(z) = log(1 + e^z) the relation softplus^(n)(z) = sigma^(n-1)(z) for n >= 1 lets us reuse the same machinery; the only change is an off-by-one in the requested order.

Sigmoid evaluation strategy

The base value s = sigma(z) is evaluated with the **framework-native** torch.sigmoid. The JAX twin uses jax.nn.sigmoid (omnibias.jax._fastpath.jax_sigmoid). Both are the stable two-branch form where(z >= 0, 1/(1+exp(-z)), exp(z)/(1+exp(z))), but they are *not* guaranteed algebraically corresponding in the extreme tails (|z| \gtrsim 40 in float64): each vendor's exp / fused kernel can differ by a few ULPs once s is near 0 or 1. Shared polynomial coefficients still make the *tower* algebraically corresponding **given the same** s; cross-backend parity tests therefore compare the closed-form derivatives on moderate z, while tests/test_sigmoid_tail_parity.py locks the documented tail contract (finite, monotone, and within a few ULPs of each other at z = \pm 40, \pm 80). Replacing either call with a hand-rolled common formula would risk breaking existing bit-identity fixtures that pin the native kernels, so we document rather than unify.

Numerical scope: shared coefficients and formulas do not establish universal floating-point identity across backends, devices, compiler transforms, or reduction orders. Compare the chosen path with dtype-appropriate tolerances; even float64 is not a general bitwise-equality guarantee. Algebraic initialization identities likewise need a numerical check for the selected coefficients and reduction order.

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.

Imports & aliases​

These names share a definition with another public module. Follow a name for its full contract, parameters, and examples.

Import from this moduleCanonical reference
sigmoid_polynomial_coeffsomnibias.core.sigmoid_polynomial_coeffs
FUNCTION

sigmoid_nth_derivative​

omnibias.torch.fastpath.eulerian.sigmoid_nth_derivativeSource ↗
sigmoid_nth_derivative(z: Tensor, n: int) -> Tensor

sigma^(n)(z) evaluated in closed form via Eulerian polynomials.

One framework-native torch.sigmoid call regardless of n, then O(n) multiply-adds in Horner form. See the module docstring for the torch / JAX tail-parity contract.

Parameters

NameType & defaultContract
zpositional-or-keywordTensorRequiredSee the operation contract and implementation below.
npositional-or-keywordintRequiredSee the operation contract and implementation below.
ReturnsTensor
Explicit validation & failure cases
  • When n < 0, raises ValueError(f'order n must be >= 0, got {n}').

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

Executed regression · every registered activation can be constructed

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.

def test_every_registered_activation_can_be_constructed(name: str) -> None:
    ombu = OperatorMultiBiasUnit(num_channels=2, K=2, base=name)
    z = torch.randn(3, 2)
    out = ombu(z)
    assert out.shape == (3, 2)
    assert torch.isfinite(out).all()
python -m pytest "packages/omnibias-torch/tests/test_unit.py::test_every_registered_activation_can_be_constructed[soft_step]" -q
Regression example · closed form tower tracks native sigmoid at tails

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.

def test_closed_form_tower_tracks_native_sigmoid_at_tails(z: float, n: int) -> None:
    """Given each backend's own ``s``, the shared Eulerian tower stays coherent."""
    t = torch_sigma_n(torch.tensor(z, dtype=torch.float64), n)
    j = jax_sigma_n(jnp.asarray(z, dtype=jnp.float64), n)
    t_v = float(t)
    j_v = float(j)
    assert math.isfinite(t_v) and math.isfinite(j_v)
    # n=0 is exactly the native sigmoid; higher orders vanish at |z|->inf.
    assert math.isclose(t_v, j_v, rel_tol=_REL_TOL, abs_tol=max(_ABS_TOL, 1e-30))
python -m pytest tests/test_sigmoid_tail_parity.py -k test_closed_form_tower_tracks_native_sigmoid_at_tails -q
Read the implementation
def sigmoid_nth_derivative(z: Tensor, n: int) -> Tensor:
    """``sigma^(n)(z)`` evaluated in closed form via Eulerian polynomials.

    One framework-native ``torch.sigmoid`` call regardless of ``n``, then
    ``O(n)`` multiply-adds in Horner form. See the module docstring for the
    torch / JAX tail-parity contract.
    """
    if n < 0:
        raise ValueError(f"order n must be >= 0, got {n}")
    # Framework-native stable sigmoid (see module docstring: not bit-identical
    # to jax.nn.sigmoid in extreme tails; polynomials are shared either way).
    s = torch.sigmoid(z)
    if n == 0:
        return s
    coeffs = sigmoid_polynomial_coeffs(n)
    return _horner(coeffs, s)
FUNCTION

sigmoid_tower​

omnibias.torch.fastpath.eulerian.sigmoid_towerSource ↗
sigmoid_tower(z: Tensor, max_order: int) -> Tensor

All derivatives through max_order from one native sigmoid evaluation.

Parameters

NameType & defaultContract
zpositional-or-keywordTensorRequiredSee the operation contract and implementation below.
max_orderpositional-or-keywordintRequiredSee the operation contract and implementation below.
ReturnsTensor
Explicit validation & failure cases
  • When max_order < 0, raises ValueError('max_order must be nonnegative').

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

Executed regression · jetmlp activation sweep matches autograd

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.

def test_jetmlp_activation_sweep_matches_autograd(base: str) -> None:
    net = JetMLP(2, 12, 1, depth=2, base=base).double()
    x = torch.randn(5, 2, dtype=torch.float64)
    h = net.hessian(x).squeeze(-1)
    h_ad = torch.func.vmap(torch.func.hessian(_val_fn(net)))(x)
    assert torch.allclose(h, h_ad, atol=1e-9)
python -m pytest "packages/omnibias-torch/tests/test_jet_pinn.py::test_jetmlp_activation_sweep_matches_autograd[sigmoid]" -q
Read the implementation
def sigmoid_tower(z: Tensor, max_order: int) -> Tensor:
    """All derivatives through max_order from one native sigmoid evaluation."""
    if max_order < 0:
        raise ValueError("max_order must be nonnegative")
    value = torch.sigmoid(z)
    rows = [value] + [_horner(sigmoid_polynomial_coeffs(n), value) for n in range(1, max_order + 1)]
    return torch.stack(rows, dim=0)
FUNCTION

softplus_nth_derivative​

omnibias.torch.fastpath.eulerian.softplus_nth_derivativeSource ↗
softplus_nth_derivative(z: Tensor, n: int) -> Tensor

softplus^(n)(z).

For n = 0 returns softplus(z) directly; for n >= 1 returns sigma^(n-1)(z).

Parameters

NameType & defaultContract
zpositional-or-keywordTensorRequiredSee the operation contract and implementation below.
npositional-or-keywordintRequiredSee the operation contract and implementation below.
ReturnsTensor
Explicit validation & failure cases
  • When n < 0, raises ValueError(f'order n must be >= 0, got {n}').

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

Executed regression · every registered activation can be constructed

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.

def test_every_registered_activation_can_be_constructed(name: str) -> None:
    ombu = OperatorMultiBiasUnit(num_channels=2, K=2, base=name)
    z = torch.randn(3, 2)
    out = ombu(z)
    assert out.shape == (3, 2)
    assert torch.isfinite(out).all()
python -m pytest "packages/omnibias-torch/tests/test_unit.py::test_every_registered_activation_can_be_constructed[soft_relu]" -q
Read the implementation
def softplus_nth_derivative(z: Tensor, n: int) -> Tensor:
    """``softplus^(n)(z)``.

    For ``n = 0`` returns ``softplus(z)`` directly; for ``n >= 1``
    returns ``sigma^(n-1)(z)``.
    """
    if n < 0:
        raise ValueError(f"order n must be >= 0, got {n}")
    if n == 0:
        return torch.nn.functional.softplus(z)
    return sigmoid_nth_derivative(z, n - 1)
FUNCTION

softplus_tower​

omnibias.torch.fastpath.eulerian.softplus_towerSource ↗
softplus_tower(z: Tensor, max_order: int) -> Tensor

Softplus value and its sigmoid derivative tower with one sigmoid call.

Parameters

NameType & defaultContract
zpositional-or-keywordTensorRequiredSee the operation contract and implementation below.
max_orderpositional-or-keywordintRequiredSee the operation contract and implementation below.
ReturnsTensor
Explicit validation & failure cases
  • When max_order < 0, raises ValueError('max_order must be nonnegative').

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

Executed regression · jetmlp activation sweep matches autograd

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.

def test_jetmlp_activation_sweep_matches_autograd(base: str) -> None:
    net = JetMLP(2, 12, 1, depth=2, base=base).double()
    x = torch.randn(5, 2, dtype=torch.float64)
    h = net.hessian(x).squeeze(-1)
    h_ad = torch.func.vmap(torch.func.hessian(_val_fn(net)))(x)
    assert torch.allclose(h, h_ad, atol=1e-9)
python -m pytest "packages/omnibias-torch/tests/test_jet_pinn.py::test_jetmlp_activation_sweep_matches_autograd[softplus]" -q
Read the implementation
def softplus_tower(z: Tensor, max_order: int) -> Tensor:
    """Softplus value and its sigmoid derivative tower with one sigmoid call."""
    if max_order < 0:
        raise ValueError("max_order must be nonnegative")
    rows = [torch.nn.functional.softplus(z)]
    if max_order:
        derivatives = sigmoid_tower(z, max_order - 1)
        rows.extend(derivatives[n] for n in range(max_order))
    return torch.stack(rows, dim=0)