omnibias.jax.activations

0.5.0rc1Release candidatePython >=3.10Apache-2.0

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

import omnibias.jax.activations

Module contract & mathematical conventions

JAX-backed activation registry.

Bit-identical mirror of the PyTorch activation dictionary. Families:

Smooth Riccati-class (closed-form fastpath for every n): sigmoid, tanh, softplus, gaussian, exp, and the trig / hyperbolic set.
Analytic products, **exact all orders**: silu, gelu, mish.
Proximal-class: arctan, log1pu2 (n in {0, 1, 2}).
Piecewise, **almost-everywhere / regular-part all-orders** tower (singular delta dropped): relu, huber, leaky_relu, prelu, relu6, hardtanh, hardsigmoid, hardswish, elu, selu, celu, softshrink, hardshrink, threshold, abs, sign, step, softsign. See omnibias.jax._activations_piecewise.
Beta-tempered smooth surrogates (-> hard as beta -> inf): soft_relu, soft_step, soft_sign (= smooth_sign), soft_abs (= softabs). See omnibias.jax._activations_tempered.

Every spec exposes forward(z), derivative(z) (== fastpath at n=1), and fastpath(z, n) (closed form up to the maximum supported 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.

Module source ↗Source revision 9a2d73df0573Running reference examples →

CONSTANT

ABS

omnibias.jax.activations.ABSSource ↗

ABS = register_activation(JaxActivationSpec(name='abs', forward=_abs_forward, derivative=_abs_derivative, fastpath=_abs_fastpath, integral=_abs_integral, riccati_polynomial=None, noise_model='none', operator_role='Absolute value; derivative sign(z) (sign(0)=0), n>=2 -> 0 (a.e.).'))

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

ABS = register_activation(
    JaxActivationSpec(
        name="abs",
        forward=_abs_forward,
        derivative=_abs_derivative,
        fastpath=_abs_fastpath,
        integral=_abs_integral,
        riccati_polynomial=None,
        noise_model="none",
        operator_role="Absolute value; derivative sign(z) (sign(0)=0), n>=2 -> 0 (a.e.).",
    )
)

CONSTANT

CELU

omnibias.jax.activations.CELUSource ↗

CELU = register_activation(make_celu_spec(_DEFAULT_CELU_ALPHA))

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

CELU = register_activation(make_celu_spec(_DEFAULT_CELU_ALPHA))

CONSTANT

ELU

omnibias.jax.activations.ELUSource ↗

ELU = register_activation(make_elu_spec(_DEFAULT_ELU_ALPHA))

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

ELU = register_activation(make_elu_spec(_DEFAULT_ELU_ALPHA))

CONSTANT

HARDSHRINK

omnibias.jax.activations.HARDSHRINKSource ↗

HARDSHRINK = register_activation(make_hardshrink_spec())

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

HARDSHRINK = register_activation(make_hardshrink_spec())

CONSTANT

HARDSIGMOID

omnibias.jax.activations.HARDSIGMOIDSource ↗

HARDSIGMOID = register_activation(JaxActivationSpec(name='hardsigmoid', forward=_hardsigmoid_forward, derivative=_hardsigmoid_derivative, fastpath=_hardsigmoid_fastpath, riccati_polynomial=None, noise_model='none', operator_role='Hard sigmoid: clamp(z/6 + 1/2, 0, 1); a.e. regular-part tower (n>=2 -> 0).', limit_pos_inf=1.0, limit_neg_inf=0.0))

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

HARDSIGMOID = register_activation(
    JaxActivationSpec(
        name="hardsigmoid",
        forward=_hardsigmoid_forward,
        derivative=_hardsigmoid_derivative,
        fastpath=_hardsigmoid_fastpath,
        riccati_polynomial=None,
        noise_model="none",
        operator_role=(
            "Hard sigmoid: clamp(z/6 + 1/2, 0, 1); a.e. regular-part tower (n>=2 -> 0)."
        ),
        limit_pos_inf=1.0,
        limit_neg_inf=0.0,
    )
)

CONSTANT

HARDSWISH

omnibias.jax.activations.HARDSWISHSource ↗

HARDSWISH = register_activation(JaxActivationSpec(name='hardswish', forward=_hardswish_forward, derivative=_hardswish_derivative, fastpath=_hardswish_fastpath, riccati_polynomial=None, noise_model='none', operator_role='Hard swish: z * hardsigmoid(z); a.e. regular-part tower (quadratic middle gives n=2 -> 1/3, n>=3 -> 0).', limit_neg_inf=0.0))

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

HARDSWISH = register_activation(
    JaxActivationSpec(
        name="hardswish",
        forward=_hardswish_forward,
        derivative=_hardswish_derivative,
        fastpath=_hardswish_fastpath,
        riccati_polynomial=None,
        noise_model="none",
        operator_role=(
            "Hard swish: z * hardsigmoid(z); a.e. regular-part tower "
            "(quadratic middle gives n=2 -> 1/3, n>=3 -> 0)."
        ),
        limit_neg_inf=0.0,
    )
)

CONSTANT

HARDTANH

omnibias.jax.activations.HARDTANHSource ↗

HARDTANH = register_activation(make_hardtanh_spec())

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

HARDTANH = register_activation(make_hardtanh_spec())

CONSTANT

LEAKY_RELU

omnibias.jax.activations.LEAKY_RELUSource ↗

LEAKY_RELU = register_activation(make_leaky_relu_spec(_DEFAULT_LEAKY_SLOPE))

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

LEAKY_RELU = register_activation(make_leaky_relu_spec(_DEFAULT_LEAKY_SLOPE))

CONSTANT

PRELU

omnibias.jax.activations.PRELUSource ↗

PRELU = register_activation(make_leaky_relu_spec(_DEFAULT_PRELU_INIT, name='prelu'))

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

PRELU = register_activation(make_leaky_relu_spec(_DEFAULT_PRELU_INIT, name="prelu"))

CONSTANT

RELU6

omnibias.jax.activations.RELU6Source ↗

RELU6 = register_activation(JaxActivationSpec(name='relu6', forward=_relu6_forward, derivative=_relu6_derivative, fastpath=_relu6_fastpath, riccati_polynomial=None, noise_model='none', operator_role='ReLU6: clamp(z, 0, 6); a.e. regular-part tower (n>=2 -> 0).', limit_pos_inf=6.0, limit_neg_inf=0.0))

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

RELU6 = register_activation(
    JaxActivationSpec(
        name="relu6",
        forward=_relu6_forward,
        derivative=_relu6_derivative,
        fastpath=_relu6_fastpath,
        riccati_polynomial=None,
        noise_model="none",
        operator_role="ReLU6: clamp(z, 0, 6); a.e. regular-part tower (n>=2 -> 0).",
        limit_pos_inf=6.0,
        limit_neg_inf=0.0,
    )
)

CONSTANT

SELU

omnibias.jax.activations.SELUSource ↗

SELU = register_activation(JaxActivationSpec(name='selu', forward=_selu_forward, derivative=_selu_derivative, fastpath=_selu_fastpath, riccati_polynomial=None, noise_model='none', operator_role='SELU: scale * ELU_{alpha}; self-normalizing; exact per-piece tower.', limit_neg_inf=-_SELU_SCALE * _SELU_ALPHA))

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

SELU = register_activation(
    JaxActivationSpec(
        name="selu",
        forward=_selu_forward,
        derivative=_selu_derivative,
        fastpath=_selu_fastpath,
        riccati_polynomial=None,
        noise_model="none",
        operator_role="SELU: scale * ELU_{alpha}; self-normalizing; exact per-piece tower.",
        limit_neg_inf=-_SELU_SCALE * _SELU_ALPHA,
    )
)

CONSTANT

SIGN

omnibias.jax.activations.SIGNSource ↗

SIGN = register_activation(JaxActivationSpec(name='sign', forward=_sign_forward, derivative=_sign_derivative, fastpath=_sign_fastpath, riccati_polynomial=None, noise_model='none', operator_role='Sign (sign(0)=0); a.e. tower is zero for all n>=1 (the singular delta at 0 is dropped). Smooth twin: soft_sign (tanh(beta z)).', limit_pos_inf=1.0, limit_neg_inf=-1.0))

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

SIGN = register_activation(
    JaxActivationSpec(
        name="sign",
        forward=_sign_forward,
        derivative=_sign_derivative,
        fastpath=_sign_fastpath,
        riccati_polynomial=None,
        noise_model="none",
        operator_role=(
            "Sign (sign(0)=0); a.e. tower is zero for all n>=1 (the singular "
            "delta at 0 is dropped). Smooth twin: soft_sign (tanh(beta z))."
        ),
        limit_pos_inf=1.0,
        limit_neg_inf=-1.0,
    )
)

CONSTANT

SOFTSHRINK

omnibias.jax.activations.SOFTSHRINKSource ↗

SOFTSHRINK = register_activation(make_softshrink_spec())

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

SOFTSHRINK = register_activation(make_softshrink_spec())

CONSTANT

SOFTSIGN

omnibias.jax.activations.SOFTSIGNSource ↗

SOFTSIGN = register_activation(JaxActivationSpec(name='softsign', forward=_softsign_forward, derivative=_softsign_derivative, fastpath=_softsign_fastpath, riccati_polynomial=None, noise_model='none', operator_role='Softsign: z/(1+|z|); exact rational per-arm tower.', limit_pos_inf=1.0, limit_neg_inf=-1.0))

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

SOFTSIGN = register_activation(
    JaxActivationSpec(
        name="softsign",
        forward=_softsign_forward,
        derivative=_softsign_derivative,
        fastpath=_softsign_fastpath,
        riccati_polynomial=None,
        noise_model="none",
        operator_role="Softsign: z/(1+|z|); exact rational per-arm tower.",
        limit_pos_inf=1.0,
        limit_neg_inf=-1.0,
    )
)

CONSTANT

STEP

omnibias.jax.activations.STEPSource ↗

STEP = register_activation(JaxActivationSpec(name='step', forward=_step_forward, derivative=_sign_derivative, fastpath=_step_fastpath, riccati_polynomial=None, noise_model='none', operator_role='Heaviside step (H(0)=0); a.e. tower is zero for all n>=1. Smooth twin: soft_step (sigmoid(beta z)).', aliases=('heaviside',), limit_pos_inf=1.0, limit_neg_inf=0.0))

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

STEP = register_activation(
    JaxActivationSpec(
        name="step",
        forward=_step_forward,
        derivative=_sign_derivative,
        fastpath=_step_fastpath,
        riccati_polynomial=None,
        noise_model="none",
        operator_role=(
            "Heaviside step (H(0)=0); a.e. tower is zero for all n>=1. Smooth "
            "twin: soft_step (sigmoid(beta z))."
        ),
        aliases=("heaviside",),
        limit_pos_inf=1.0,
        limit_neg_inf=0.0,
    )
)

CONSTANT

THRESHOLD

omnibias.jax.activations.THRESHOLDSource ↗

THRESHOLD = register_activation(make_threshold_spec())

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

THRESHOLD = register_activation(make_threshold_spec())

FUNCTION

make_celu_spec

omnibias.jax.activations.make_celu_specSource ↗

make_celu_spec(alpha: float=_DEFAULT_CELU_ALPHA, *, name: str='celu') -> JaxActivationSpec

CELU spec with fixed alpha (exact on each open piece, all orders).

Parameters

Name	Type & default	Contract
`alpha`positional-or-keyword	`float`Default: _DEFAULT_CELU_ALPHA	See the operation contract and implementation below.
`name`keyword-only	`str`Default: 'celu'	See the operation contract and implementation below.

ReturnsJaxActivationSpec

Explicit validation & failure cases

When alpha == 0, raises ValueError('CELU alpha must be nonzero.').

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

Read the implementation

def make_celu_spec(alpha: float = _DEFAULT_CELU_ALPHA, *, name: str = "celu") -> JaxActivationSpec:
    """CELU spec with fixed ``alpha`` (exact on each open piece, all orders)."""
    if alpha == 0:
        raise ValueError("CELU alpha must be nonzero.")

    def fwd(z: Array) -> Array:
        return _celu_forward(z, alpha)

    def deriv(z: Array) -> Array:
        return _celu_derivative(z, alpha)

    def fp(z: Array, n: int) -> Array:
        return _celu_fastpath(z, n, alpha)

    return JaxActivationSpec(
        name=name,
        forward=fwd,
        derivative=deriv,
        fastpath=fp,
        riccati_polynomial=None,
        noise_model="none",
        operator_role="CELU: z for z>0, alpha*(exp(z/alpha)-1) otherwise; exact per-piece tower.",
        limit_neg_inf=-alpha,
    )

FUNCTION

make_elu_spec

omnibias.jax.activations.make_elu_specSource ↗

make_elu_spec(alpha: float=_DEFAULT_ELU_ALPHA, *, name: str='elu') -> JaxActivationSpec

ELU spec with fixed alpha (exact on each open piece, all orders).

Parameters

Name	Type & default	Contract
`alpha`positional-or-keyword	`float`Default: _DEFAULT_ELU_ALPHA	See the operation contract and implementation below.
`name`keyword-only	`str`Default: 'elu'	See the operation contract and implementation below.

ReturnsJaxActivationSpec

Read the implementation

def make_elu_spec(alpha: float = _DEFAULT_ELU_ALPHA, *, name: str = "elu") -> JaxActivationSpec:
    """ELU spec with fixed ``alpha`` (exact on each open piece, all orders)."""

    def fwd(z: Array) -> Array:
        return _elu_forward(z, alpha)

    def deriv(z: Array) -> Array:
        return _elu_derivative(z, alpha)

    def integ(z: Array) -> Array:
        return _elu_integral(z, alpha)

    def fp(z: Array, n: int) -> Array:
        return _elu_fastpath(z, n, alpha)

    return JaxActivationSpec(
        name=name,
        forward=fwd,
        derivative=deriv,
        fastpath=fp,
        integral=integ,
        riccati_polynomial=None,
        noise_model="none",
        operator_role="ELU: z for z>0, alpha*(exp(z)-1) otherwise; exact per-piece tower.",
        limit_neg_inf=-alpha,
    )

FUNCTION

make_hardshrink_spec

omnibias.jax.activations.make_hardshrink_specSource ↗

make_hardshrink_spec(lambd: float=_DEFAULT_SHRINK_LAMBD, *, name: str='hardshrink') -> JaxActivationSpec

Hard-shrinkage spec with threshold lambd (discontinuous at +/-lambd).

Parameters

Name	Type & default	Contract
`lambd`positional-or-keyword	`float`Default: _DEFAULT_SHRINK_LAMBD	See the operation contract and implementation below.
`name`keyword-only	`str`Default: 'hardshrink'	See the operation contract and implementation below.

ReturnsJaxActivationSpec

Read the implementation

def make_hardshrink_spec(
    lambd: float = _DEFAULT_SHRINK_LAMBD, *, name: str = "hardshrink"
) -> JaxActivationSpec:
    """Hard-shrinkage spec with threshold ``lambd`` (discontinuous at +/-lambd)."""

    def fwd(z: Array) -> Array:
        return _hardshrink_forward(z, lambd)

    def deriv(z: Array) -> Array:
        return _hardshrink_derivative(z, lambd)

    def fp(z: Array, n: int) -> Array:
        return _ae_linear_fastpath(z, n, fwd, deriv)

    return JaxActivationSpec(
        name=name,
        forward=fwd,
        derivative=deriv,
        fastpath=fp,
        riccati_polynomial=None,
        noise_model="none",
        operator_role="Hard-shrink; a.e. regular-part tower (n>=2 -> 0).",
    )

FUNCTION

make_hardtanh_spec

omnibias.jax.activations.make_hardtanh_specSource ↗

make_hardtanh_spec(min_val: float=_DEFAULT_HARDTANH_MIN, max_val: float=_DEFAULT_HARDTANH_MAX, *, name: str='hardtanh') -> JaxActivationSpec

Hardtanh spec clamping to [min_val, max_val] (a.e. all-orders tower).

Parameters

Name	Type & default	Contract
`min_val`positional-or-keyword	`float`Default: _DEFAULT_HARDTANH_MIN	See the operation contract and implementation below.
`max_val`positional-or-keyword	`float`Default: _DEFAULT_HARDTANH_MAX	See the operation contract and implementation below.
`name`keyword-only	`str`Default: 'hardtanh'	See the operation contract and implementation below.

ReturnsJaxActivationSpec

Read the implementation

def make_hardtanh_spec(
    min_val: float = _DEFAULT_HARDTANH_MIN,
    max_val: float = _DEFAULT_HARDTANH_MAX,
    *,
    name: str = "hardtanh",
) -> JaxActivationSpec:
    """Hardtanh spec clamping to ``[min_val, max_val]`` (a.e. all-orders tower)."""

    def fwd(z: Array) -> Array:
        return _hardtanh_forward(z, min_val, max_val)

    def deriv(z: Array) -> Array:
        return _hardtanh_derivative(z, min_val, max_val)

    def fp(z: Array, n: int) -> Array:
        return _ae_linear_fastpath(z, n, fwd, deriv)

    return JaxActivationSpec(
        name=name,
        forward=fwd,
        derivative=deriv,
        fastpath=fp,
        riccati_polynomial=None,
        noise_model="none",
        operator_role="Hardtanh: clamp(z, min, max); a.e. regular-part tower (n>=2 -> 0).",
        limit_pos_inf=max_val,
        limit_neg_inf=min_val,
    )

FUNCTION

make_leaky_relu_spec

omnibias.jax.activations.make_leaky_relu_specSource ↗

make_leaky_relu_spec(slope: float=_DEFAULT_LEAKY_SLOPE, *, name: str='leaky_relu') -> JaxActivationSpec

Leaky-ReLU spec with fixed negative slope (a.e. all-orders tower).

Parameters

Name	Type & default	Contract
`slope`positional-or-keyword	`float`Default: _DEFAULT_LEAKY_SLOPE	See the operation contract and implementation below.
`name`keyword-only	`str`Default: 'leaky_relu'	See the operation contract and implementation below.

ReturnsJaxActivationSpec

Read the implementation

def make_leaky_relu_spec(
    slope: float = _DEFAULT_LEAKY_SLOPE, *, name: str = "leaky_relu"
) -> JaxActivationSpec:
    """Leaky-ReLU spec with fixed negative ``slope`` (a.e. all-orders tower)."""

    def fwd(z: Array) -> Array:
        return _leaky_relu_forward(z, slope)

    def deriv(z: Array) -> Array:
        return _leaky_relu_derivative(z, slope)

    def integ(z: Array) -> Array:
        return _leaky_relu_integral(z, slope)

    def fp(z: Array, n: int) -> Array:
        return _ae_linear_fastpath(z, n, fwd, deriv)

    return JaxActivationSpec(
        name=name,
        forward=fwd,
        derivative=deriv,
        fastpath=fp,
        integral=integ,
        riccati_polynomial=None,
        noise_model="none",
        operator_role=(
            "Leaky ReLU: z for z>0, slope*z otherwise; a.e. regular-part tower "
            "(n>=2 -> 0)."
        ),
    )

FUNCTION

make_softshrink_spec

omnibias.jax.activations.make_softshrink_specSource ↗

make_softshrink_spec(lambd: float=_DEFAULT_SHRINK_LAMBD, *, name: str='softshrink') -> JaxActivationSpec

Soft-shrinkage spec (proximal of L1) with threshold lambd.

Parameters

Name	Type & default	Contract
`lambd`positional-or-keyword	`float`Default: _DEFAULT_SHRINK_LAMBD	See the operation contract and implementation below.
`name`keyword-only	`str`Default: 'softshrink'	See the operation contract and implementation below.

ReturnsJaxActivationSpec

Read the implementation

def make_softshrink_spec(
    lambd: float = _DEFAULT_SHRINK_LAMBD, *, name: str = "softshrink"
) -> JaxActivationSpec:
    """Soft-shrinkage spec (proximal of L1) with threshold ``lambd``."""

    def fwd(z: Array) -> Array:
        return _softshrink_forward(z, lambd)

    def deriv(z: Array) -> Array:
        return _softshrink_derivative(z, lambd)

    def fp(z: Array, n: int) -> Array:
        return _ae_linear_fastpath(z, n, fwd, deriv)

    return JaxActivationSpec(
        name=name,
        forward=fwd,
        derivative=deriv,
        fastpath=fp,
        riccati_polynomial=None,
        noise_model="none",
        operator_role="Soft-shrink (L1 proximal); a.e. regular-part tower (n>=2 -> 0).",
    )

FUNCTION

make_threshold_spec

omnibias.jax.activations.make_threshold_specSource ↗

make_threshold_spec(threshold: float=_DEFAULT_THRESHOLD, value: float=_DEFAULT_THRESHOLD_VALUE, *, name: str='threshold') -> JaxActivationSpec

Threshold spec: identity above threshold, constant value below.

Parameters

Name	Type & default	Contract
`threshold`positional-or-keyword	`float`Default: _DEFAULT_THRESHOLD	See the operation contract and implementation below.
`value`positional-or-keyword	`float`Default: _DEFAULT_THRESHOLD_VALUE	See the operation contract and implementation below.
`name`keyword-only	`str`Default: 'threshold'	See the operation contract and implementation below.

ReturnsJaxActivationSpec

Read the implementation

def make_threshold_spec(
    threshold: float = _DEFAULT_THRESHOLD,
    value: float = _DEFAULT_THRESHOLD_VALUE,
    *,
    name: str = "threshold",
) -> JaxActivationSpec:
    """Threshold spec: identity above ``threshold``, constant ``value`` below."""

    def fwd(z: Array) -> Array:
        return _threshold_forward(z, threshold, value)

    def deriv(z: Array) -> Array:
        return _threshold_derivative(z, threshold)

    def fp(z: Array, n: int) -> Array:
        return _ae_linear_fastpath(z, n, fwd, deriv)

    return JaxActivationSpec(
        name=name,
        forward=fwd,
        derivative=deriv,
        fastpath=fp,
        riccati_polynomial=None,
        noise_model="none",
        operator_role="Threshold gate; a.e. regular-part tower (n>=2 -> 0).",
        limit_neg_inf=value,
    )

CONSTANT

SOFT_RELU

omnibias.jax.activations.SOFT_RELUSource ↗

SOFT_RELU = register_activation(make_soft_relu_spec())

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

SOFT_RELU = register_activation(make_soft_relu_spec())

CONSTANT

SOFT_STEP

omnibias.jax.activations.SOFT_STEPSource ↗

SOFT_STEP = register_activation(make_soft_step_spec())

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

SOFT_STEP = register_activation(make_soft_step_spec())

FUNCTION

make_soft_leaky_relu_spec

omnibias.jax.activations.make_soft_leaky_relu_specSource ↗

make_soft_leaky_relu_spec(negative_slope: float=0.01, beta: float=_DEFAULT_BETA, *, name: str='soft_leaky_relu') -> JaxActivationSpec

Smooth leaky ReLU alpha*z + (1-alpha)*soft_relu(z; beta) (all orders).

Parameters

Name	Type & default	Contract
`negative_slope`positional-or-keyword	`float`Default: 0.01	See the operation contract and implementation below.
`beta`positional-or-keyword	`float`Default: _DEFAULT_BETA	See the operation contract and implementation below.
`name`keyword-only	`str`Default: 'soft_leaky_relu'	See the operation contract and implementation below.

ReturnsJaxActivationSpec

Executed regression · soft leaky relu converges

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_soft_leaky_relu_converges() -> None:
    alpha = 0.1
    z = jnp.asarray([-3.0, -1.0, -0.3, 0.3, 1.0, 3.0])
    hard = jnp.where(z > 0, z, alpha * z)
    err = float(jnp.max(jnp.abs(make_soft_leaky_relu_spec(alpha, 128.0).forward(z) - hard)))
    assert err < 1e-2

python -m pytest "packages/omnibias-jax/tests/test_tempered.py::test_soft_leaky_relu_converges" -q

Read the implementation

def make_soft_leaky_relu_spec(
    negative_slope: float = 0.01,
    beta: float = _DEFAULT_BETA,
    *,
    name: str = "soft_leaky_relu",
) -> JaxActivationSpec:
    """Smooth leaky ReLU ``alpha*z + (1-alpha)*soft_relu(z; beta)`` (all orders)."""
    alpha = negative_slope
    one_minus_alpha = 1.0 - alpha

    def _soft_relu_nth(z: Array, n: int) -> Array:
        return (beta ** (n - 1)) * _fastpath.softplus_nth_derivative(beta * z, n)

    def fwd(z: Array) -> Array:
        return alpha * z + one_minus_alpha * (
            jnp.logaddexp(jnp.zeros_like(z), beta * z) / beta
        )

    def fp(z: Array, n: int) -> Array:
        if n < 0:
            raise ValueError(f"order n must be >= 0, got {n}.")
        if n == 0:
            return fwd(z)
        soft = one_minus_alpha * _soft_relu_nth(z, n)
        if n == 1:
            return jnp.full_like(z, alpha) + soft
        return soft

    def deriv(z: Array) -> Array:
        return fp(z, 1)

    return JaxActivationSpec(
        name=name,
        forward=fwd,
        derivative=deriv,
        fastpath=fp,
        riccati_polynomial=None,
        noise_model="none",
        operator_role=(
            "Smooth leaky ReLU: alpha*z + (1-alpha)*softplus(beta z)/beta "
            "-> leaky_relu as beta -> inf."
        ),
    )

FUNCTION

make_soft_relu_spec

omnibias.jax.activations.make_soft_relu_specSource ↗

make_soft_relu_spec(beta: float=_DEFAULT_BETA, *, name: str='soft_relu') -> JaxActivationSpec

Smooth ReLU surrogate softplus(beta z) / beta (all-orders tower).

Parameters

Name	Type & default	Contract
`beta`positional-or-keyword	`float`Default: _DEFAULT_BETA	See the operation contract and implementation below.
`name`keyword-only	`str`Default: 'soft_relu'	See the operation contract and implementation below.

ReturnsJaxActivationSpec

Executed regression · soft relu converges to relu

Full file, imports & fixtures ↗

def test_soft_relu_converges_to_relu() -> None:
    relu = get_activation("relu")
    z = jnp.asarray([-3.0, -1.0, -0.3, 0.3, 1.0, 3.0])
    prev = None
    for beta in (2.0, 8.0, 32.0, 128.0):
        err = float(jnp.max(jnp.abs(make_soft_relu_spec(beta).forward(z) - relu.forward(z))))
        if prev is not None:
            assert err < prev
        prev = err
    assert prev < 1e-2

python -m pytest "packages/omnibias-jax/tests/test_tempered.py::test_soft_relu_converges_to_relu" -q

Regression example · soft relu beta scaling identity

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_soft_relu_beta_scaling_identity(beta: float) -> None:
    spec = make_soft_relu_spec(beta)
    softplus = get_activation("softplus")
    z = jnp.linspace(-3.0, 3.0, 41)
    for n in range(0, 6):
        lhs = np.asarray(spec.fastpath(z, n))
        rhs = np.asarray((beta ** (n - 1)) * softplus.fastpath(beta * z, n))
        np.testing.assert_allclose(lhs, rhs, rtol=1e-6, atol=1e-6, err_msg=f"n={n}")

python -m pytest packages/omnibias-jax/tests/test_tempered.py -k test_soft_relu_beta_scaling_identity -q

Read the implementation

def make_soft_relu_spec(
    beta: float = _DEFAULT_BETA, *, name: str = "soft_relu"
) -> JaxActivationSpec:
    """Smooth ReLU surrogate ``softplus(beta z) / beta`` (all-orders tower)."""
    return tempered(
        SOFTPLUS,
        beta,
        scale="one_over_beta",
        name=name,
        operator_role=(
            "Smooth ReLU: softplus(beta z)/beta -> relu as beta -> inf; "
            "the n=2 bump -> delta."
        ),
        limit_neg_inf=0.0,
    )

FUNCTION

make_soft_step_spec

omnibias.jax.activations.make_soft_step_specSource ↗

make_soft_step_spec(beta: float=_DEFAULT_BETA, *, name: str='soft_step') -> JaxActivationSpec

Smooth step surrogate sigmoid(beta z) (all-orders tower).

Parameters

Name	Type & default	Contract
`beta`positional-or-keyword	`float`Default: _DEFAULT_BETA	See the operation contract and implementation below.
`name`keyword-only	`str`Default: 'soft_step'	See the operation contract and implementation below.

ReturnsJaxActivationSpec

Read the implementation

def make_soft_step_spec(
    beta: float = _DEFAULT_BETA, *, name: str = "soft_step"
) -> JaxActivationSpec:
    """Smooth step surrogate ``sigmoid(beta z)`` (all-orders tower)."""
    return tempered(
        SIGMOID,
        beta,
        scale="unit",
        name=name,
        aliases=("soft_heaviside",) if name == "soft_step" else (),
        operator_role="Smooth Heaviside: sigmoid(beta z) -> step as beta -> inf.",
        limit_pos_inf=1.0,
        limit_neg_inf=0.0,
    )

FUNCTION

make_swish_spec

omnibias.jax.activations.make_swish_specSource ↗

make_swish_spec(beta: float=_DEFAULT_BETA, *, name: str='swish') -> JaxActivationSpec

Swish z * sigmoid(beta z) with tunable beta (silu is beta=1).

Parameters

Name	Type & default	Contract
`beta`positional-or-keyword	`float`Default: _DEFAULT_BETA	See the operation contract and implementation below.
`name`keyword-only	`str`Default: 'swish'	See the operation contract and implementation below.

ReturnsJaxActivationSpec

Executed regression · swish beta1 is silu

Full file, imports & fixtures ↗

def test_swish_beta1_is_silu() -> None:
    swish = make_swish_spec(1.0)
    silu = get_activation("silu")
    z = jnp.linspace(-4.0, 4.0, 41)
    for n in range(0, 6):
        np.testing.assert_allclose(
            np.asarray(swish.fastpath(z, n)), np.asarray(silu.fastpath(z, n)),
            rtol=1e-6, atol=1e-6, err_msg=f"n={n}",
        )

python -m pytest "packages/omnibias-jax/tests/test_tempered.py::test_swish_beta1_is_silu" -q

Regression example · swish matches autograd

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_swish_matches_autograd() -> None:
    spec = make_swish_spec(1.5)
    z = jnp.linspace(-4.0, 4.0, 41)
    fwd = lambda x: spec.forward(x)  # noqa: E731
    for n in range(1, 5):
        closed = np.asarray(spec.fastpath(z, n))
        auto = np.asarray(jax.vmap(_nth_grad(fwd, n))(z))
        np.testing.assert_allclose(closed, auto, rtol=1e-6, atol=1e-6, err_msg=f"n={n}")

python -m pytest packages/omnibias-jax/tests/test_tempered.py -k test_swish_matches_autograd -q

Read the implementation

def make_swish_spec(beta: float = _DEFAULT_BETA, *, name: str = "swish") -> JaxActivationSpec:
    """Swish ``z * sigmoid(beta z)`` with tunable ``beta`` (``silu`` is beta=1)."""

    def fwd(z: Array) -> Array:
        return z * _fastpath.jax_sigmoid(beta * z)

    def fp(z: Array, n: int) -> Array:
        if n < 0:
            raise ValueError(f"order n must be >= 0, got {n}.")
        if n == 0:
            return fwd(z)
        bz = beta * z
        term_a = z * (beta**n) * _fastpath.sigmoid_nth_derivative(bz, n)
        term_b = n * (beta ** (n - 1)) * _fastpath.sigmoid_nth_derivative(bz, n - 1)
        return term_a + term_b

    def deriv(z: Array) -> Array:
        return fp(z, 1)

    return JaxActivationSpec(
        name=name,
        forward=fwd,
        derivative=deriv,
        fastpath=fp,
        riccati_polynomial=None,
        noise_model="none",
        operator_role="Swish: z * sigmoid(beta z); exact all-orders (silu is beta=1).",
    )

FUNCTION

tempered_activation

omnibias.jax.activations.tempered_activationSource ↗

tempered_activation(base: str | JaxActivationSpec, beta: float | Array=_DEFAULT_BETA, *, scale: str='one_over_beta', name: str | None=None) -> JaxActivationSpec

Functional learnable-temperature surrogate (JAX twin of the module blocks).

Resolve base (registry name or spec) and return a tempered JaxActivationSpec whose whole closed-form tower is scaled by beta -- which may be a traced Array (e.g. a learnable temperature threaded through jax.grad). Mirrors the torch/keras TemperedActivation layer without holding state.

Parameters

Name	Type & default	Contract
`base`positional-or-keyword	`str \| JaxActivationSpec`Required	See the operation contract and implementation below.
`beta`positional-or-keyword	`float \| Array`Default: _DEFAULT_BETA	See the operation contract and implementation below.
`scale`keyword-only	`str`Default: 'one_over_beta'	See the operation contract and implementation below.
`name`keyword-only	`str \| None`Default: None	See the operation contract and implementation below.

ReturnsJaxActivationSpec

Executed regression · tempered activation jit traced beta

Full file, imports & fixtures ↗

def test_tempered_activation_jit_traced_beta() -> None:
    z = jnp.linspace(-3.0, 3.0, 41)

    def bump(beta):
        return tempered_activation("softplus", beta=beta).fastpath(z, 2)

    jitted = jax.jit(bump)(3.0)
    np.testing.assert_allclose(np.asarray(jitted), np.asarray(bump(3.0)), rtol=1e-6, atol=1e-6)

python -m pytest "packages/omnibias-jax/tests/test_tempered.py::test_tempered_activation_jit_traced_beta" -q

Regression example · tempered activation traced beta gradient

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_tempered_activation_traced_beta_gradient() -> None:
    z = jnp.linspace(-3.0, 3.0, 41)

    def loss(beta):
        return tempered_activation("softplus", beta=beta, scale="one_over_beta").forward(z).sum()

    g = float(jax.grad(loss)(2.0))
    eps = 1e-6
    fd = (float(loss(2.0 + eps)) - float(loss(2.0 - eps))) / (2 * eps)
    assert abs(g - fd) < 1e-4, f"grad {g} vs finite-diff {fd}"

python -m pytest packages/omnibias-jax/tests/test_tempered.py -k test_tempered_activation_traced_beta_gradient -q

Read the implementation

def tempered_activation(
    base: str | JaxActivationSpec,
    beta: float | Array = _DEFAULT_BETA,
    *,
    scale: str = "one_over_beta",
    name: str | None = None,
) -> JaxActivationSpec:
    """Functional learnable-temperature surrogate (JAX twin of the module blocks).

    Resolve ``base`` (registry name or spec) and return a tempered
    :class:`JaxActivationSpec` whose whole closed-form tower is scaled by
    ``beta`` -- which may be a traced ``Array`` (e.g. a learnable temperature
    threaded through ``jax.grad``). Mirrors the torch/keras
    ``TemperedActivation`` layer without holding state.
    """
    spec = get_activation(base)
    return tempered(
        spec,
        beta,
        scale=scale,
        name=name if name is not None else f"tempered_{spec.name}",
    )

CONSTANT

ARCTAN

omnibias.jax.activations.ARCTANSource ↗

ARCTAN = register_activation(JaxActivationSpec(name='arctan', forward=_arctan_forward, derivative=_arctan_derivative, fastpath=_arctan_fastpath, integral=_arctan_integral, riccati_polynomial=None, noise_model='cauchy', operator_role='K=2 collapse -> 1/(1+z^2); Cauchy IRLS weight.', aliases=('atan',), limit_pos_inf=math.pi / 2.0, limit_neg_inf=-math.pi / 2.0))

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

ARCTAN = register_activation(
    JaxActivationSpec(
        name="arctan",
        forward=_arctan_forward,
        derivative=_arctan_derivative,
        fastpath=_arctan_fastpath,
        integral=_arctan_integral,
        riccati_polynomial=None,
        noise_model="cauchy",
        operator_role="K=2 collapse -> 1/(1+z^2); Cauchy IRLS weight.",
        aliases=("atan",),
        limit_pos_inf=math.pi / 2.0,
        limit_neg_inf=-math.pi / 2.0,
    )
)

CONSTANT

COS

omnibias.jax.activations.COSSource ↗

COS = register_activation(JaxActivationSpec(name='cos', forward=_cos_forward, derivative=_cos_derivative, fastpath=_cos_fastpath, integral=_cos_integral, riccati_polynomial=None, noise_model='none', operator_role='Plane-wave / Bloch embedding; cos(z + n*pi/2) gives every order in closed form.'))

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

COS = register_activation(
    JaxActivationSpec(
        name="cos",
        forward=_cos_forward,
        derivative=_cos_derivative,
        fastpath=_cos_fastpath,
        integral=_cos_integral,
        riccati_polynomial=None,
        noise_model="none",
        operator_role=(
            "Plane-wave / Bloch embedding; cos(z + n*pi/2) gives every order in closed form."
        ),
    )
)

CONSTANT

COSH

omnibias.jax.activations.COSHSource ↗

COSH = register_activation(JaxActivationSpec(name='cosh', forward=_cosh_forward, derivative=_cosh_derivative, fastpath=_cosh_fastpath, integral=_cosh_integral, riccati_polynomial=None, noise_model='none', operator_role='Hyperbolic even primitive; eigenfunction of d^2/dz^2 with eigenvalue +1.'))

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

COSH = register_activation(
    JaxActivationSpec(
        name="cosh",
        forward=_cosh_forward,
        derivative=_cosh_derivative,
        fastpath=_cosh_fastpath,
        integral=_cosh_integral,
        riccati_polynomial=None,
        noise_model="none",
        operator_role=("Hyperbolic even primitive; eigenfunction of d^2/dz^2 with eigenvalue +1."),
    )
)

CONSTANT

COT

omnibias.jax.activations.COTSource ↗

COT = register_activation(JaxActivationSpec(name='cot', forward=_cot_forward, derivative=_cot_derivative, fastpath=_cot_fastpath, integral=_cot_integral, riccati_polynomial=(-1.0, 0.0, -1.0), noise_model='none', operator_role="Riccati-class periodic activation; cot'(z) = -(1 + cot^2(z)).", aliases=('ctg', 'ctan')))

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

COT = register_activation(
    JaxActivationSpec(
        name="cot",
        forward=_cot_forward,
        derivative=_cot_derivative,
        fastpath=_cot_fastpath,
        integral=_cot_integral,
        riccati_polynomial=(-1.0, 0.0, -1.0),
        noise_model="none",
        operator_role=("Riccati-class periodic activation; cot'(z) = -(1 + cot^2(z))."),
        aliases=("ctg", "ctan"),
    )
)

CONSTANT

COTH

omnibias.jax.activations.COTHSource ↗

COTH = register_activation(JaxActivationSpec(name='coth', forward=_coth_forward, derivative=_coth_derivative, fastpath=_coth_fastpath, integral=_coth_integral, riccati_polynomial=(1.0, 0.0, -1.0), noise_model='none', operator_role="Hyperbolic Riccati: coth'(z) = 1 - coth^2(z); same polynomial shape as tanh, but coth has a pole at z=0."))

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

COTH = register_activation(
    JaxActivationSpec(
        name="coth",
        forward=_coth_forward,
        derivative=_coth_derivative,
        fastpath=_coth_fastpath,
        integral=_coth_integral,
        riccati_polynomial=(1.0, 0.0, -1.0),
        noise_model="none",
        operator_role=(
            "Hyperbolic Riccati: coth'(z) = 1 - coth^2(z); same polynomial "
            "shape as tanh, but coth has a pole at z=0."
        ),
    )
)

CONSTANT

EXP

omnibias.jax.activations.EXPSource ↗

EXP = register_activation(JaxActivationSpec(name='exp', forward=_exp_forward, derivative=_exp_forward, fastpath=_exp_fastpath, integral=_exp_forward, riccati_polynomial=(0.0, 1.0), noise_model='poisson', operator_role='K=2 collapse -> exp(z) (eigenfunction of d/dz); log-link Newton step for Poisson regression.', limit_neg_inf=0.0))

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

EXP = register_activation(
    JaxActivationSpec(
        name="exp",
        forward=_exp_forward,
        derivative=_exp_forward,
        fastpath=_exp_fastpath,
        integral=_exp_forward,
        riccati_polynomial=(0.0, 1.0),
        noise_model="poisson",
        operator_role=(
            "K=2 collapse -> exp(z) (eigenfunction of d/dz); "
            "log-link Newton step for Poisson regression."
        ),
        limit_neg_inf=0.0,  # exp -> 0 as z -> -inf; diverges as z -> +inf
    )
)

CONSTANT

GAUSSIAN

omnibias.jax.activations.GAUSSIANSource ↗

GAUSSIAN = register_activation(JaxActivationSpec(name='gaussian', forward=_gaussian, derivative=_gaussian_derivative, fastpath=_fastpath.gaussian_nth_derivative, integral=_gaussian_integral, riccati_polynomial=None, noise_model='gaussian_kernel', operator_role='K=2 collapse -> -z * exp(-z^2/2); RBF/Hermite spectral basis, Laplacian-of-Gaussian conv kernel.', aliases=('rbf',), limit_pos_inf=0.0, limit_neg_inf=0.0))

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

GAUSSIAN = register_activation(
    JaxActivationSpec(
        name="gaussian",
        forward=_gaussian,
        derivative=_gaussian_derivative,
        fastpath=_fastpath.gaussian_nth_derivative,
        integral=_gaussian_integral,
        riccati_polynomial=None,
        noise_model="gaussian_kernel",
        operator_role=(
            "K=2 collapse -> -z * exp(-z^2/2); RBF/Hermite spectral basis, "
            "Laplacian-of-Gaussian conv kernel."
        ),
        aliases=("rbf",),
        limit_pos_inf=0.0,
        limit_neg_inf=0.0,
    )
)

CONSTANT

GELU

omnibias.jax.activations.GELUSource ↗

GELU = register_activation(JaxActivationSpec(name='gelu', forward=_gelu_forward, derivative=_gelu_derivative, fastpath=_gelu_fastpath, integral=_gelu_integral, riccati_polynomial=None, noise_model='none', operator_role='K=2 collapse -> Phi(z) + z * phi(z); smoothed half-space gate.'))

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

GELU = register_activation(
    JaxActivationSpec(
        name="gelu",
        forward=_gelu_forward,
        derivative=_gelu_derivative,
        fastpath=_gelu_fastpath,
        integral=_gelu_integral,
        riccati_polynomial=None,
        noise_model="none",
        operator_role="K=2 collapse -> Phi(z) + z * phi(z); smoothed half-space gate.",
    )
)

CONSTANT

HUBER

omnibias.jax.activations.HUBERSource ↗

HUBER = register_activation(JaxActivationSpec(name='huber', forward=_huber_forward, derivative=_huber_derivative, fastpath=_huber_fastpath, integral=_huber_integral, riccati_polynomial=None, noise_model='laplace_smoothed', operator_role='K=2 collapse -> clip(z, -tau, tau); ISTA soft-shrink step.'))

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

HUBER = register_activation(
    JaxActivationSpec(
        name="huber",
        forward=_huber_forward,
        derivative=_huber_derivative,
        fastpath=_huber_fastpath,
        integral=_huber_integral,
        riccati_polynomial=None,
        noise_model="laplace_smoothed",
        operator_role=("K=2 collapse -> clip(z, -tau, tau); ISTA soft-shrink step."),
    )
)

CONSTANT

IntegralFn

omnibias.jax.activations.IntegralFnSource ↗

IntegralFn = Callable[[Array], Array]

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

IntegralFn = Callable[[Array], Array]

CONSTANT

JaxActivationSpec

omnibias.jax.JaxActivationSpecSource ↗

JaxActivationSpec: TypeAlias = _CoreActivationSpec[Array]

JAX-specialised omnibias.core.spec.ActivationSpec. The field shape is identical to the torch backend's spec; TensorT is pinned to jax.Array. Re-exported under the JaxActivationSpec name to preserve the v0.1 public API.

Read the implementation

JaxActivationSpec: TypeAlias = _CoreActivationSpec[Array]

Other public import paths

omnibias.jax.JaxActivationSpec
omnibias.jax.activations.JaxActivationSpec

CONSTANT

LOG1PU2

omnibias.jax.activations.LOG1PU2Source ↗

LOG1PU2 = register_activation(JaxActivationSpec(name='log1pu2', forward=_log1pu2_forward, derivative=_log1pu2_derivative, fastpath=_log1pu2_fastpath, integral=_log1pu2_integral, riccati_polynomial=None, noise_model='redescending_m', operator_role='K=2 collapse -> 2 z / (1 + z^2); redescending M-estimator.', aliases=('logcosh_dual',)))

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

LOG1PU2 = register_activation(
    JaxActivationSpec(
        name="log1pu2",
        forward=_log1pu2_forward,
        derivative=_log1pu2_derivative,
        fastpath=_log1pu2_fastpath,
        integral=_log1pu2_integral,
        riccati_polynomial=None,
        noise_model="redescending_m",
        operator_role=("K=2 collapse -> 2 z / (1 + z^2); redescending M-estimator."),
        aliases=("logcosh_dual",),
    )
)

CONSTANT

LOG_COSH

omnibias.jax.activations.LOG_COSHSource ↗

LOG_COSH = register_activation(JaxActivationSpec(name='log_cosh', forward=_log_cosh_forward, derivative=_log_cosh_derivative, fastpath=_log_cosh_fastpath, riccati_polynomial=None, noise_model='laplace_smoothed', operator_role='Smooth-|z| activation: log(cosh(z)) ≈ |z| for large |z|, ≈ z²/2 for small. Standard log-amplitude in spin-NQS; first derivative is tanh, second is 1-tanh².', aliases=('logcosh',)))

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

LOG_COSH = register_activation(
    JaxActivationSpec(
        name="log_cosh",
        forward=_log_cosh_forward,
        derivative=_log_cosh_derivative,
        fastpath=_log_cosh_fastpath,
        riccati_polynomial=None,
        noise_model="laplace_smoothed",
        operator_role=(
            "Smooth-|z| activation: log(cosh(z)) ≈ |z| for large |z|, "
            "≈ z²/2 for small. Standard log-amplitude in spin-NQS; "
            "first derivative is tanh, second is 1-tanh²."
        ),
        aliases=("logcosh",),
    )
)

CONSTANT

MISH

omnibias.jax.activations.MISHSource ↗

MISH = register_activation(JaxActivationSpec(name='mish', forward=_mish_forward, derivative=_mish_derivative, fastpath=_mish_fastpath, riccati_polynomial=None, noise_model='none', operator_role='Mish: z * tanh(softplus(z)). Self-gated residual activation useful in transformer FFN blocks; smoother than ReLU.'))

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

MISH = register_activation(
    JaxActivationSpec(
        name="mish",
        forward=_mish_forward,
        derivative=_mish_derivative,
        fastpath=_mish_fastpath,
        riccati_polynomial=None,
        noise_model="none",
        operator_role=(
            "Mish: z * tanh(softplus(z)). Self-gated residual activation "
            "useful in transformer FFN blocks; smoother than ReLU."
        ),
    )
)

CONSTANT

RELU

omnibias.jax.activations.RELUSource ↗

RELU = register_activation(JaxActivationSpec(name='relu', forward=_relu_forward, derivative=_relu_derivative, fastpath=_relu_fastpath, integral=_relu_integral, riccati_polynomial=None, noise_model='none', operator_role='K=2 collapse -> Heaviside step; equality-constraint indicator.'))

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

RELU = register_activation(
    JaxActivationSpec(
        name="relu",
        forward=_relu_forward,
        derivative=_relu_derivative,
        fastpath=_relu_fastpath,
        integral=_relu_integral,
        riccati_polynomial=None,
        noise_model="none",
        operator_role=("K=2 collapse -> Heaviside step; equality-constraint indicator."),
    )
)

CONSTANT

SECH

omnibias.jax.activations.SECHSource ↗

SECH = register_activation(JaxActivationSpec(name='sech', forward=_sech_forward, derivative=_sech_derivative, fastpath=_sech_fastpath, integral=_sech_integral, riccati_polynomial=None, noise_model='none', operator_role='Soliton bound-state amplitude (Poschl-Teller, sine-Gordon, KdV, NLS soliton profile).'))

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

SECH = register_activation(
    JaxActivationSpec(
        name="sech",
        forward=_sech_forward,
        derivative=_sech_derivative,
        fastpath=_sech_fastpath,
        integral=_sech_integral,
        riccati_polynomial=None,
        noise_model="none",
        operator_role=(
            "Soliton bound-state amplitude (Poschl-Teller, sine-Gordon, KdV, NLS soliton profile)."
        ),
    )
)

CONSTANT

SIGMOID

omnibias.jax.activations.SIGMOIDSource ↗

SIGMOID = register_activation(JaxActivationSpec(name='sigmoid', forward=_sigmoid, derivative=_sigmoid_derivative, fastpath=_fastpath.sigmoid_nth_derivative, tower=_fastpath.sigmoid_tower, integral=_sigmoid_integral, riccati_polynomial=(0.0, 1.0, -1.0), noise_model='bernoulli', operator_role='K=2 collapse -> Bernoulli variance s(1-s); Newton/IRLS step for logistic regression.', aliases=('logistic',), limit_pos_inf=1.0, limit_neg_inf=0.0))

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

SIGMOID = register_activation(
    JaxActivationSpec(
        name="sigmoid",
        forward=_sigmoid,
        derivative=_sigmoid_derivative,
        fastpath=_fastpath.sigmoid_nth_derivative,
        tower=_fastpath.sigmoid_tower,
        integral=_sigmoid_integral,
        riccati_polynomial=(0.0, 1.0, -1.0),
        noise_model="bernoulli",
        operator_role=(
            "K=2 collapse -> Bernoulli variance s(1-s); Newton/IRLS step for logistic regression."
        ),
        aliases=("logistic",),
        limit_pos_inf=1.0,
        limit_neg_inf=0.0,
    )
)

CONSTANT

SILU

omnibias.jax.activations.SILUSource ↗

SILU = register_activation(JaxActivationSpec(name='silu', forward=_silu_forward, derivative=_silu_derivative, fastpath=_silu_fastpath, riccati_polynomial=None, noise_model='none', operator_role='K=2 collapse -> sigmoid + z * sigmoid * (1 - sigmoid); smoothed gate for transformer FFN compatibility.', aliases=('swish',)))

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

SILU = register_activation(
    JaxActivationSpec(
        name="silu",
        forward=_silu_forward,
        derivative=_silu_derivative,
        fastpath=_silu_fastpath,
        riccati_polynomial=None,
        noise_model="none",
        operator_role=(
            "K=2 collapse -> sigmoid + z * sigmoid * (1 - sigmoid); "
            "smoothed gate for transformer FFN compatibility."
        ),
        aliases=("swish",),
    )
)

CONSTANT

SIN

omnibias.jax.activations.SINSource ↗

SIN = register_activation(JaxActivationSpec(name='sin', forward=_sin_forward, derivative=_sin_derivative, fastpath=_sin_fastpath, integral=_sin_integral, riccati_polynomial=None, noise_model='none', operator_role='Plane-wave / Bloch embedding; sin(z + n*pi/2) gives every order in closed form.'))

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

SIN = register_activation(
    JaxActivationSpec(
        name="sin",
        forward=_sin_forward,
        derivative=_sin_derivative,
        fastpath=_sin_fastpath,
        integral=_sin_integral,
        riccati_polynomial=None,
        noise_model="none",
        operator_role=(
            "Plane-wave / Bloch embedding; sin(z + n*pi/2) gives every order in closed form."
        ),
    )
)

CONSTANT

SINH

omnibias.jax.activations.SINHSource ↗

SINH = register_activation(JaxActivationSpec(name='sinh', forward=_sinh_forward, derivative=_sinh_derivative, fastpath=_sinh_fastpath, integral=_sinh_integral, riccati_polynomial=None, noise_model='none', operator_role='Hyperbolic odd primitive; eigenfunction of d^2/dz^2 with eigenvalue +1.'))

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

SINH = register_activation(
    JaxActivationSpec(
        name="sinh",
        forward=_sinh_forward,
        derivative=_sinh_derivative,
        fastpath=_sinh_fastpath,
        integral=_sinh_integral,
        riccati_polynomial=None,
        noise_model="none",
        operator_role=("Hyperbolic odd primitive; eigenfunction of d^2/dz^2 with eigenvalue +1."),
    )
)

CONSTANT

SMOOTH_SIGN

omnibias.jax.activations.SMOOTH_SIGNSource ↗

SMOOTH_SIGN = register_activation(JaxActivationSpec(name='smooth_sign', forward=_smooth_sign_forward, derivative=_smooth_sign_derivative, fastpath=_smooth_sign_fastpath, integral=_smooth_sign_integral, riccati_polynomial=None, noise_model='symmetric_bernoulli_tempered', operator_role='Temperature-controlled smooth sign: tanh(z/T) → sign(z) as T → 0. The beta-tempered sign surrogate; used in variational annealing schedules.', aliases=('soft_sign',), limit_pos_inf=1.0, limit_neg_inf=-1.0))

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

SMOOTH_SIGN = register_activation(
    JaxActivationSpec(
        name="smooth_sign",
        forward=_smooth_sign_forward,
        derivative=_smooth_sign_derivative,
        fastpath=_smooth_sign_fastpath,
        integral=_smooth_sign_integral,
        riccati_polynomial=None,
        noise_model="symmetric_bernoulli_tempered",
        operator_role=(
            "Temperature-controlled smooth sign: tanh(z/T) → sign(z) as "
            "T → 0. The beta-tempered sign surrogate; used in variational "
            "annealing schedules."
        ),
        aliases=("soft_sign",),
        limit_pos_inf=1.0,
        limit_neg_inf=-1.0,
    )
)

CONSTANT

SOFTABS

omnibias.jax.activations.SOFTABSSource ↗

SOFTABS = register_activation(JaxActivationSpec(name='softabs', forward=_softabs_forward, derivative=_softabs_derivative, fastpath=_softabs_fastpath, integral=_softabs_integral, riccati_polynomial=None, noise_model='huber_smoothed', operator_role='Smooth absolute value: sqrt(z² + ε²) - ε. C² everywhere; the eps-tempered abs surrogate (-> |z| as eps -> 0). Useful in Jastrow factors and complex-amplitude magnitudes.', aliases=('soft_abs',)))

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

SOFTABS = register_activation(
    JaxActivationSpec(
        name="softabs",
        forward=_softabs_forward,
        derivative=_softabs_derivative,
        fastpath=_softabs_fastpath,
        integral=_softabs_integral,
        riccati_polynomial=None,
        noise_model="huber_smoothed",
        operator_role=(
            "Smooth absolute value: sqrt(z² + ε²) - ε. C² everywhere; the "
            "eps-tempered abs surrogate (-> |z| as eps -> 0). Useful in "
            "Jastrow factors and complex-amplitude magnitudes."
        ),
        aliases=("soft_abs",),
    )
)

CONSTANT

SOFTPLUS

omnibias.jax.activations.SOFTPLUSSource ↗

SOFTPLUS = register_activation(JaxActivationSpec(name='softplus', forward=_softplus, derivative=_softplus_derivative, fastpath=_fastpath.softplus_nth_derivative, tower=_fastpath.softplus_tower, riccati_polynomial=None, noise_model='bernoulli', operator_role='K=2 collapse -> sigmoid (Bernoulli mean); canonical PINN base for diffusion-class PDEs.', limit_neg_inf=0.0))

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

SOFTPLUS = register_activation(
    JaxActivationSpec(
        name="softplus",
        forward=_softplus,
        derivative=_softplus_derivative,
        fastpath=_fastpath.softplus_nth_derivative,
        tower=_fastpath.softplus_tower,
        riccati_polynomial=None,
        noise_model="bernoulli",
        operator_role=(
            "K=2 collapse -> sigmoid (Bernoulli mean); "
            "canonical PINN base for diffusion-class PDEs."
        ),
        limit_neg_inf=0.0,  # softplus -> 0 as z -> -inf; diverges as z -> +inf
    )
)

CONSTANT

TAN

omnibias.jax.activations.TANSource ↗

TAN = register_activation(JaxActivationSpec(name='tan', forward=_tan_forward, derivative=_tan_derivative, fastpath=_tan_fastpath, integral=_tan_integral, riccati_polynomial=(1.0, 0.0, 1.0), noise_model='none', operator_role="Riccati-class periodic activation; tan'(z) = 1 + tan^2(z)."))

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

TAN = register_activation(
    JaxActivationSpec(
        name="tan",
        forward=_tan_forward,
        derivative=_tan_derivative,
        fastpath=_tan_fastpath,
        integral=_tan_integral,
        riccati_polynomial=(1.0, 0.0, 1.0),
        noise_model="none",
        operator_role=("Riccati-class periodic activation; tan'(z) = 1 + tan^2(z)."),
    )
)

CONSTANT

TANH

omnibias.jax.activations.TANHSource ↗

TANH = register_activation(JaxActivationSpec(name='tanh', forward=_tanh, derivative=_tanh_derivative, fastpath=_fastpath.tanh_nth_derivative, tower=_fastpath.tanh_tower, integral=_tanh_integral, riccati_polynomial=(1.0, 0.0, -1.0), noise_model='symmetric_bernoulli', operator_role='K=2 collapse -> 1 - tanh^2; symmetric IRLS bell.', limit_pos_inf=1.0, limit_neg_inf=-1.0))

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

TANH = register_activation(
    JaxActivationSpec(
        name="tanh",
        forward=_tanh,
        derivative=_tanh_derivative,
        fastpath=_fastpath.tanh_nth_derivative,
        tower=_fastpath.tanh_tower,
        integral=_tanh_integral,
        riccati_polynomial=(1.0, 0.0, -1.0),
        noise_model="symmetric_bernoulli",
        operator_role="K=2 collapse -> 1 - tanh^2; symmetric IRLS bell.",
        limit_pos_inf=1.0,
        limit_neg_inf=-1.0,
    )
)

FUNCTION

get_activation

omnibias.jax.get_activationSource ↗

get_activation(name: str | JaxActivationSpec) -> JaxActivationSpec

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
`name`positional-or-keyword	`str \| JaxActivationSpec`Required	See the operation contract and implementation below.

ReturnsJaxActivationSpec

Explicit validation & failure cases

When not isinstance(name, str), raises TypeError(f'Expected str or JaxActivationSpec, got {type(name).__name__}: {name!r}').
When key not in _REGISTRY, raises KeyError(f'Unknown activation {name!r}. Known: {known}.').

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

Executed regression · soft sign soft abs aliases

Full file, imports & fixtures ↗

def test_soft_sign_soft_abs_aliases() -> None:
    assert get_activation("soft_sign").name == "smooth_sign"
    assert get_activation("soft_abs").name == "softabs"

python -m pytest "packages/omnibias-jax/tests/test_tempered.py::test_soft_sign_soft_abs_aliases" -q

Regression example · negative order parity jax

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_negative_order_parity_jax(name: str) -> None:
    spec = jax_get(name)
    assert spec.fastpath is not None
    with pytest.raises(ValueError, match="order n must be"):
        spec.fastpath(jnp.zeros(1), -1)

python -m pytest tests/test_jax_parity_invariants.py -k test_negative_order_parity_jax -q

Read the implementation

def get_activation(name: str | JaxActivationSpec) -> JaxActivationSpec:
    if isinstance(name, _CoreActivationSpec):
        return name
    if not isinstance(name, str):
        raise TypeError(f"Expected str or JaxActivationSpec, got {type(name).__name__}: {name!r}")
    key = name.lower()
    if key not in _REGISTRY:
        known = ", ".join(sorted({s.name for s in _REGISTRY.values()}))
        raise KeyError(f"Unknown activation {name!r}. Known: {known}.")
    return _REGISTRY[key]

Other public import paths

omnibias.jax.get_activation
omnibias.jax.activations.get_activation

FUNCTION

is_registered

omnibias.jax.is_registeredSource ↗

is_registered(name: str) -> bool

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

Returnsbool

Read the implementation

def is_registered(name: str) -> bool:
    return isinstance(name, str) and name.lower() in _REGISTRY

Other public import paths

omnibias.jax.is_registered
omnibias.jax.activations.is_registered

FUNCTION

list_activations

omnibias.jax.list_activationsSource ↗

list_activations() -> list[str]

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

Returnslist[str]

Executed regression · all registered

Full file, imports & fixtures ↗

def test_all_registered() -> None:
    registered = set(list_activations())
    for name in _PIECEWISE_NAMES:
        assert name in registered

python -m pytest "packages/omnibias-jax/tests/test_piecewise.py::test_all_registered" -q

Regression example · list activations matches torch

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_list_activations_matches_torch() -> None:
    """Both backends register the same set of activation names."""
    from omnibias.torch.activations.registry import list_activations as torch_list

    assert sorted(torch_list()) == sorted(jax_list())

python -m pytest tests/test_jax_parity.py -k test_list_activations_matches_torch -q

Read the implementation

def list_activations() -> list[str]:
    return sorted({s.name for s in _REGISTRY.values()})

Other public import paths

omnibias.jax.list_activations
omnibias.jax.activations.list_activations

FUNCTION

register_activation

omnibias.jax.register_activationSource ↗

register_activation(spec: JaxActivationSpec) -> JaxActivationSpec

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

ReturnsJaxActivationSpec

Explicit validation & failure cases

When not spec.name, raises ValueError('JaxActivationSpec.name must be a non-empty string.').
When name in _REGISTRY and _REGISTRY[name] is not spec, raises ValueError(f'Activation {name!r} is already registered with a different spec.').
When akey in _REGISTRY and _REGISTRY[akey] is not spec, raises ValueError(f'Alias {alias!r} for activation {name!r} clashes.').

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

Read the implementation

def register_activation(spec: JaxActivationSpec) -> JaxActivationSpec:
    if not spec.name:
        raise ValueError("JaxActivationSpec.name must be a non-empty string.")
    name = spec.name.lower()
    if name in _REGISTRY and _REGISTRY[name] is not spec:
        raise ValueError(f"Activation {name!r} is already registered with a different spec.")
    _REGISTRY[name] = spec
    for alias in spec.aliases:
        akey = alias.lower()
        if akey in _REGISTRY and _REGISTRY[akey] is not spec:
            raise ValueError(f"Alias {alias!r} for activation {name!r} clashes.")
        _REGISTRY[akey] = spec
    return spec

Other public import paths

omnibias.jax.register_activation
omnibias.jax.activations.register_activation

← Package overview Runnable learning examples →

omnibias.jax.activations

ABS​

CELU​

ELU​

HARDSHRINK​

HARDSIGMOID​

HARDSWISH​

HARDTANH​

LEAKY_RELU​

PRELU​

RELU6​

SELU​

SIGN​

SOFTSHRINK​

SOFTSIGN​

STEP​

THRESHOLD​

make_celu_spec​

Parameters

make_elu_spec​

Parameters

make_hardshrink_spec​

Parameters

make_hardtanh_spec​

Parameters

make_leaky_relu_spec​

Parameters

make_softshrink_spec​

Parameters

make_threshold_spec​

Parameters

SOFT_RELU​

SOFT_STEP​

make_soft_leaky_relu_spec​

Parameters

make_soft_relu_spec​

Parameters

make_soft_step_spec​

Parameters

make_swish_spec​

Parameters

tempered_activation​

Parameters

ARCTAN​

COS​

COSH​

COT​

COTH​

EXP​

GAUSSIAN​

GELU​

HUBER​

IntegralFn​

JaxActivationSpec​

LOG1PU2​

LOG_COSH​

MISH​

RELU​

SECH​

SIGMOID​

SILU​

SIN​

SINH​

SMOOTH_SIGN​

SOFTABS​

SOFTPLUS​

TAN​

TANH​

get_activation​

Parameters

is_registered​

Parameters

list_activations​

register_activation​

Parameters

ABS

CELU

ELU

HARDSHRINK

HARDSIGMOID

HARDSWISH

HARDTANH

LEAKY_RELU

PRELU

RELU6

SELU

SIGN

SOFTSHRINK

SOFTSIGN

STEP

THRESHOLD

make_celu_spec

make_elu_spec

make_hardshrink_spec

make_hardtanh_spec

make_leaky_relu_spec

make_softshrink_spec

make_threshold_spec

SOFT_RELU

SOFT_STEP

make_soft_leaky_relu_spec

make_soft_relu_spec

make_soft_step_spec

make_swish_spec

tempered_activation

ARCTAN

COS

COSH

COT

COTH

EXP

GAUSSIAN

GELU

HUBER

IntegralFn

JaxActivationSpec

LOG1PU2

LOG_COSH

MISH

RELU

SECH

SIGMOID

SILU

SIN

SINH

SMOOTH_SIGN

SOFTABS

SOFTPLUS

TAN

TANH

get_activation

is_registered

list_activations

register_activation