omnibias.keras.blocks.operator

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.blocks.operator

Module contract & mathematical conventions

Operator-typed wrapper around the Keras OperatorMultiBiasUnit.

An OperatorBlock carries an operator role tag -- "identity", "grad", "laplacian", "derivative", "band", or "integral" -- that selects both the multi-bias arity K and the forward behaviour. Mirrors omnibias.torch.blocks.operator.OperatorBlock.

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

OpName

omnibias.keras.OpNameSource ↗

OpName = Literal['identity', 'grad', 'laplacian', 'derivative', 'band', 'integral']

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

OpName = Literal["identity", "grad", "laplacian", "derivative", "band", "integral"]

Other public import paths

omnibias.keras.OpName
omnibias.keras.blocks.OpName
omnibias.keras.blocks.operator.OpName

CLASS

OperatorBlock

omnibias.keras.OperatorBlockSource ↗

class OperatorBlock(layers.Layer)

An OMBU with operator-type semantics (Keras backend).

Derivative-op semantics (important precondition)

The derivative ops (grad, laplacian, derivative) evaluate the closed-form **mean-bias fast-path** sigma^(order)(z + mean(biases)) (see omnibias.keras.unit.OperatorMultiBiasUnit.analytic_derivative), *not* the literal finite-difference stencil sum_k s_k sigma(z + b_k) over the individual biases. The two coincide only under the **identity-nesting** initialisation -- init_delta = 0 (the block's default for these ops), which seeds all K bias columns equal so the mean *is* the shared bias and the closed-form sigma^(order) at that point is the intended derivative.

With **untied biases** that have drifted apart during training (or a nonzero init_delta), mean(biases) no longer characterises the stencil and the op returns the single-point closed-form derivative *at the mean bias* rather than the literal multi-bias stencil. Keep init_delta = 0 (and, if exact stencil semantics matter, learnable_biases=False) for the derivative ops; the activation-level math stays mathematically equivalent to the torch / jax backends under this precondition.

Parameters

op : one of identity / grad / laplacian / derivative / band / integral base : str or ActivationSpec, default "gaussian" channels : int, default 1 init_delta : float, optional Initial bias spread. Defaults to 0 for derivative ops (the identity-nesting precondition above) and 1 for band / integral. A nonzero value breaks the mean-bias == literal-stencil equivalence for the derivative ops. learnable_biases : bool, default True learnable_signs : bool, optional derivative_order : int, optional Required for op="derivative" (block uses K = order + 1). normalize_integral : bool, default False integral_small_width : float, default 1e-4

Backend scope: this identifies the corresponding mathematical construction. It does not guarantee bitwise-equal floating-point results across devices, compilers, dtypes, or backend versions. Validate the selected path with an appropriate absolute and relative tolerance.

Parameters

Name	Type & default	Contract
`op`positional-or-keyword	`OpName`Default: 'identity'	one of `identity / grad / laplacian / derivative / band / integral`
`base`positional-or-keyword	`str \| ActivationSpec[Any]`Default: 'gaussian'	str or `ActivationSpec`, default `"gaussian"`
`channels`positional-or-keyword	`int`Default: 1	int, default 1
`init_bias`positional-or-keyword	`float`Default: 0.0	See the operation contract and implementation below.
`init_delta`positional-or-keyword	`float \| None`Default: None	float, optional Initial bias spread. Defaults to `0` for derivative ops (the identity-nesting precondition above) and `1` for `band` / `integral`. A nonzero value breaks the mean-bias == literal-stencil equivalence for the derivative ops.
`learnable_biases`positional-or-keyword	`bool`Default: True	bool, default True
`learnable_signs`positional-or-keyword	`bool \| None`Default: None	bool, optional
`derivative_order`positional-or-keyword	`int \| None`Default: None	int, optional Required for `op="derivative"` (block uses `K = order + 1`).
`normalize_integral`positional-or-keyword	`bool`Default: False	bool, default False
`integral_small_width`positional-or-keyword	`float`Default: 0.0001	float, default 1e-4
`kwargs`variadic keyword	`Any`Required	See the operation contract and implementation below.

Executed regression · invalid op raises

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_invalid_op_raises() -> None:
    with pytest.raises(ValueError):
        OperatorBlock(channels=2, op="nonsense")

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

Regression example · integral block runs

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 -k test_integral_block_runs -q

Fields & methods

__init__method

__init__(self, op: OpName='identity', base: str | ActivationSpec[Any]='gaussian', channels: int=1, init_bias: float=0.0, init_delta: float | None=None, learnable_biases: bool=True, learnable_signs: bool | None=None, derivative_order: int | None=None, normalize_integral: bool=False, integral_small_width: float=0.0001, **kwargs: Any) -> None

See the class contract above and the exact method implementation below.

Parameters

Name	Type & default	Contract
`op`positional-or-keyword	`OpName`Default: 'identity'	See the operation contract and implementation below.
`base`positional-or-keyword	`str \| ActivationSpec[Any]`Default: 'gaussian'	See the operation contract and implementation below.
`channels`positional-or-keyword	`int`Default: 1	See the operation contract and implementation below.
`init_bias`positional-or-keyword	`float`Default: 0.0	See the operation contract and implementation below.
`init_delta`positional-or-keyword	`float \| None`Default: None	See the operation contract and implementation below.
`learnable_biases`positional-or-keyword	`bool`Default: True	See the operation contract and implementation below.
`learnable_signs`positional-or-keyword	`bool \| None`Default: None	See the operation contract and implementation below.
`derivative_order`positional-or-keyword	`int \| None`Default: None	See the operation contract and implementation below.
`normalize_integral`positional-or-keyword	`bool`Default: False	See the operation contract and implementation below.
`integral_small_width`positional-or-keyword	`float`Default: 0.0001	See the operation contract and implementation below.
`kwargs`variadic keyword	`Any`Required	See the operation contract and implementation below.

Returns: None

Validation & failure cases

When op not in _VALID_OPS, raises ValueError(f'op must be one of {list(_VALID_OPS)}, got {op!r}.').
When derivative_order is not None and op in ('grad', 'laplacian'), raises ValueError(f"derivative_order is only valid with op='derivative'; use op={op!r} without it.").
When derivative_order is None, raises ValueError("OperatorBlock(op='derivative') requires derivative_order.").
When derivative_order < 1, raises ValueError(f'derivative_order must be >= 1, got {derivative_order}.').
When spec.fastpath is None, raises TypeError(f'OperatorBlock(op={op!r}) requires a base activation with a closed-form derivative kernel; activation {spec.name!r} has none.').
When op == 'integral' and spec.integral is None, raises TypeError(f"OperatorBlock(op='integral') requires a base activation with a closed-form integral kernel; activation {spec.name!r} has none.").

Executed regression · invalid op raises

Full file, imports & fixtures ↗

def test_invalid_op_raises() -> None:
    with pytest.raises(ValueError):
        OperatorBlock(channels=2, op="nonsense")

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

Method implementation

    def __init__(
        self,
        op: OpName = "identity",
        base: str | ActivationSpec[Any] = "gaussian",
        channels: int = 1,
        init_bias: float = 0.0,
        init_delta: float | None = None,
        learnable_biases: bool = True,
        learnable_signs: bool | None = None,
        derivative_order: int | None = None,
        normalize_integral: bool = False,
        integral_small_width: float = 1e-4,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        if op not in _VALID_OPS:
            raise ValueError(f"op must be one of {list(_VALID_OPS)}, got {op!r}.")
        self.op = op
        self.normalize_integral = normalize_integral
        self.integral_small_width = integral_small_width
        if derivative_order is not None and op in ("grad", "laplacian"):
            raise ValueError(
                f"derivative_order is only valid with op='derivative'; use op={op!r} without it."
            )
        if op == "derivative":
            if derivative_order is None:
                raise ValueError("OperatorBlock(op='derivative') requires derivative_order.")
            if derivative_order < 1:
                raise ValueError(f"derivative_order must be >= 1, got {derivative_order}.")
            K = derivative_order + 1
            self._derivative_order = derivative_order
        else:
            K = _OP_K[op]
            self._derivative_order = _OP_DERIVATIVE_ORDER[op]
        spec = base if isinstance(base, ActivationSpec) else get_activation(base)
        self._base_name = base if isinstance(base, str) else base.name

        if op in ("grad", "laplacian", "derivative"):
            required_order = self._derivative_order
            if spec.fastpath is None:
                raise TypeError(
                    f"OperatorBlock(op={op!r}) requires a base activation with a "
                    f"closed-form derivative kernel; activation {spec.name!r} has none."
                )
            try:
                spec.fastpath(ops.zeros((1,)), required_order)
            except NotImplementedError as e:
                raise TypeError(
                    f"OperatorBlock(op={op!r}) requires order-{required_order} "
                    f"derivative; activation {spec.name!r} only supports a partial "
                    f"fastpath. ({e})"
                ) from None
        if op == "integral" and spec.integral is None:
            raise TypeError(
                f"OperatorBlock(op='integral') requires a base activation with a "
                f"closed-form integral kernel; activation {spec.name!r} has none."
            )

        if init_delta is None:
            init_delta = 1.0 if op in ("band", "integral") else 0.0
        if learnable_signs is None:
            learnable_signs = False
        self.init_delta = init_delta
        self.learnable_biases = learnable_biases
        self.learnable_signs = learnable_signs

        dtype = self.variable_dtype
        if op == "band":
            signs_init = _window_signs(channels, K, dtype=dtype)
            offsets = central_bias_offsets(K, init_delta, dtype=dtype)[None, :]
            biases_init = identity_init_biases(channels, K, bias_value=init_bias, dtype=dtype) + offsets
        elif op == "integral":
            signs_init = _window_signs(channels, K, dtype=dtype)
            biases_init = np.empty((channels, K), dtype=dtype)
            biases_init[:, 0] = init_bias
            biases_init[:, 1] = _inverse_softplus(init_delta)
        else:
            signs_init = None
            biases_init = None

        self.ombu = OperatorMultiBiasUnit(
            num_channels=channels,
            K=K,
            base=spec,
            init_bias=init_bias,
            init_delta=init_delta if op not in ("band", "integral") else 0.0,
            learnable_biases=learnable_biases,
            learnable_signs=learnable_signs,
            biases_init=biases_init,
            signs_init=signs_init,
        )
        self.built = True

omnibias.keras.blocks.operator

OpName​

OperatorBlock​

Parameters

Fields & methods

Parameters

Parameters

Parameters

OpName

OperatorBlock