API REFERENCE/omnibias-keras

omnibias.keras.fastpath.dispatch

0.0.2a1AlphaPython >=3.10Apache-2.0

Keras 3 operator layers and activation paths implemented with keras.ops for TensorFlow, JAX, or Torch.

import omnibias.keras.fastpath.dispatch

Module contract & mathematical conventions

Forward-pass helpers shared by the OMBU and OperatorBlock layers, written against keras.ops so they run on any Keras backend.

• multibias_literal_forward: sum_k s_k * sigma(z + b_k).
• multibias_collapsed_forward: closed-form sigma^(order)(z + bias_mean).
• multibias_integral_forward: sign-weighted antiderivative sum.
• multibias_integral_window_forward: ordered bias-window integral.

Closed-form fast paths avoid the finite-difference quotient entirely, so autograd through them (on whichever Keras backend is active) is stable.

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

multibias_collapsed_forward

omnibias.keras.fastpath.multibias_collapsed_forwardSource ↗

multibias_collapsed_forward(z: Any, biases: Any, spec: ActivationSpec[Any], order: int) -> Any

Closed-form sigma^(order)(z + bias_mean).

Parameters

Name	Type & default	Contract
zpositional-or-keyword	AnyRequired	See the operation contract and implementation below.
biasespositional-or-keyword	AnyRequired	See the operation contract and implementation below.
specpositional-or-keyword	ActivationSpec[Any]Required	See the operation contract and implementation below.
orderpositional-or-keyword	intRequired	See the operation contract and implementation below.

ReturnsAny

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 {spec.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.

Read the implementation

def multibias_collapsed_forward(
    z: Any,
    biases: Any,
    spec: ActivationSpec[Any],
    order: int,
) -> Any:
    """Closed-form ``sigma^(order)(z + bias_mean)``."""
    if order < 0:
        raise ValueError(f"order must be >= 0, got {order}.")
    if spec.fastpath is None:
        raise NotImplementedError(
            f"Activation {spec.name!r} has no closed-form derivative kernel."
        )
    b_mean = ops.mean(biases, axis=-1)  # (C,)
    return spec.fastpath(z + b_mean, order)

Other public import paths

• omnibias.keras.fastpath.multibias_collapsed_forward
• omnibias.keras.fastpath.dispatch.multibias_collapsed_forward

FUNCTION

multibias_integral_forward

omnibias.keras.fastpath.multibias_integral_forwardSource ↗

multibias_integral_forward(z: Any, biases: Any, signs: Any, spec: ActivationSpec[Any]) -> Any

Closed-form sign-weighted antiderivative forward.

Parameters

Name	Type & default	Contract
zpositional-or-keyword	AnyRequired	See the operation contract and implementation below.
biasespositional-or-keyword	AnyRequired	See the operation contract and implementation below.
signspositional-or-keyword	AnyRequired	See the operation contract and implementation below.
specpositional-or-keyword	ActivationSpec[Any]Required	See the operation contract and implementation below.

ReturnsAny

Explicit validation & failure cases

• When spec.integral is None, raises NotImplementedError(f'Activation {spec.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.

Executed regression · integral block runs

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_integral_block_runs() -> None:
    z = np.random.default_rng(4).normal(size=(8, 4))
    block = OperatorBlock(channels=4, op="integral", base="gaussian", init_delta=1.0)
    out = _np(block(_t(z)))
    assert out.shape == (8, 4)
    assert np.all(np.isfinite(out))

python -m pytest "packages/omnibias-keras/tests/test_blocks.py::test_integral_block_runs" -q

Read the implementation

def multibias_integral_forward(
    z: Any,
    biases: Any,
    signs: Any,
    spec: ActivationSpec[Any],
) -> Any:
    """Closed-form sign-weighted antiderivative forward."""
    if spec.integral is None:
        raise NotImplementedError(
            f"Activation {spec.name!r} has no closed-form integral kernel."
        )
    return multibias_literal_forward(z, biases, signs, spec.integral)

Related APIsomnibias.keras.fastpath.multibias_literal_forward↗

Other public import paths

• omnibias.keras.fastpath.multibias_integral_forward
• omnibias.keras.fastpath.dispatch.multibias_integral_forward

FUNCTION

multibias_integral_window_forward

omnibias.keras.fastpath.multibias_integral_window_forwardSource ↗

multibias_integral_window_forward(z: Any, center: Any, width: Any, signs: Any, spec: ActivationSpec[Any], *, normalize: bool=False, small_width_threshold: float=0.0001, use_small_width_taylor: bool=True) -> Any

Closed-form ordered bias-window integral S(z + b_hi) - S(z + b_lo).

Parameters

Name	Type & default	Contract
zpositional-or-keyword	AnyRequired	See the operation contract and implementation below.
centerpositional-or-keyword	AnyRequired	See the operation contract and implementation below.
widthpositional-or-keyword	AnyRequired	See the operation contract and implementation below.
signspositional-or-keyword	AnyRequired	See the operation contract and implementation below.
specpositional-or-keyword	ActivationSpec[Any]Required	See the operation contract and implementation below.
normalizekeyword-only	boolDefault: False	See the operation contract and implementation below.
small_width_thresholdkeyword-only	floatDefault: 0.0001	See the operation contract and implementation below.
use_small_width_taylorkeyword-only	boolDefault: True	See the operation contract and implementation below.

ReturnsAny

Executed regression · integral block runs

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_integral_block_runs() -> None:
    z = np.random.default_rng(4).normal(size=(8, 4))
    block = OperatorBlock(channels=4, op="integral", base="gaussian", init_delta=1.0)
    out = _np(block(_t(z)))
    assert out.shape == (8, 4)
    assert np.all(np.isfinite(out))

python -m pytest "packages/omnibias-keras/tests/test_blocks.py::test_integral_block_runs" -q

Read the implementation

def multibias_integral_window_forward(
    z: Any,
    center: Any,
    width: Any,
    signs: Any,
    spec: ActivationSpec[Any],
    *,
    normalize: bool = False,
    small_width_threshold: float = 1e-4,
    use_small_width_taylor: bool = True,
) -> Any:
    """Closed-form ordered bias-window integral ``S(z + b_hi) - S(z + b_lo)``."""
    half_width = 0.5 * width
    endpoints = ops.stack((center - half_width, center + half_width), axis=-1)
    exact = multibias_integral_forward(z, endpoints, signs, spec)
    if normalize:
        exact = exact / ops.maximum(width, 1e-12)

    if not use_small_width_taylor or small_width_threshold <= 0.0:
        return exact

    midpoint = z + center
    approx = spec.forward(midpoint)
    if not normalize:
        approx = approx * width
    mask = width < small_width_threshold
    return ops.where(mask, approx, exact)

Related APIsomnibias.keras.fastpath.multibias_integral_forward↗

Other public import paths

• omnibias.keras.fastpath.multibias_integral_window_forward
• omnibias.keras.fastpath.dispatch.multibias_integral_window_forward

FUNCTION

multibias_literal_forward

omnibias.keras.fastpath.multibias_literal_forwardSource ↗

multibias_literal_forward(z: Any, biases: Any, signs: Any, sigma: Callable[[Any], Any]) -> Any

Literal multi-bias activation.

out[..., c] = sum_k signs[c, k] * sigma(z[..., c] + biases[c, k]).

Parameters

z : tensor of shape (..., C) biases : tensor of shape (C, K) signs : tensor of shape (C, K) or (K,) sigma : callable applied elementwise.

Parameters

Name	Type & default	Contract
zpositional-or-keyword	AnyRequired	tensor of shape (..., C)
biasespositional-or-keyword	AnyRequired	tensor of shape (C, K)
signspositional-or-keyword	AnyRequired	tensor of shape (C, K) or (K,)
sigmapositional-or-keyword	Callable[[Any], Any]Required	callable applied elementwise.

ReturnsAny

Executed regression · cmbconv1d shape

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_cmbconv1d_shape() -> None:
    x = np.random.default_rng(6).normal(size=(4, 10, 3))
    layer = cmbConv1D(filters=5, kernel_size=3, op="identity", base="tanh")
    out = _np(layer(_t(x)))
    assert out.shape[0] == 4 and out.shape[-1] == 5

python -m pytest "packages/omnibias-keras/tests/test_blocks.py::test_cmbconv1d_shape" -q

Read the implementation

def multibias_literal_forward(
    z: Any,
    biases: Any,
    signs: Any,
    sigma: Callable[[Any], Any],
) -> Any:
    """Literal multi-bias activation.

    ``out[..., c] = sum_k signs[c, k] * sigma(z[..., c] + biases[c, k])``.

    Parameters
    ----------
    z : tensor of shape ``(..., C)``
    biases : tensor of shape ``(C, K)``
    signs : tensor of shape ``(C, K)`` or ``(K,)``
    sigma : callable applied elementwise.
    """
    num_channels = int(biases.shape[0])
    K = int(biases.shape[1])
    signs2d = _as_2d_signs(signs, num_channels, K)
    z_expand = ops.expand_dims(z, -1)  # (..., C, 1)
    activated = sigma(z_expand + biases)  # (..., C, K)
    return ops.sum(activated * signs2d, axis=-1)

Other public import paths

• omnibias.keras.fastpath.multibias_literal_forward
• omnibias.keras.fastpath.dispatch.multibias_literal_forward

← Package overviewRunnable learning examples →