API REFERENCE/omnibias-keras

omnibias.keras.tempered_blocks

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.tempered_blocks

Module contract & mathematical conventions

Learnable-temperature activation layers (Keras backend).

keras.layers.Layer twins of omnibias.torch.tempered_blocks: they hold a differentiable temperature beta (or negative slope alpha) and call the closed-form tempered / piecewise tower each forward, so the same weight trains on TensorFlow, JAX, or PyTorch through keras.ops.

• TemperedActivation -- softplus(beta z)/beta -> relu, sigmoid(beta z) -> step, tanh(beta z) -> sign (pick the base + scale).
• LearnablePReLU -- leaky ReLU with a learnable negative slope on the almost-everywhere tower.

beta (and alpha) should stay positive; add your own reparameterisation if you need to guarantee it.

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 →

CLASS

LearnablePReLU

omnibias.keras.LearnablePReLUSource ↗

class LearnablePReLU(layers.Layer)

Leaky ReLU with a (learnable) negative slope alpha on the a.e. tower.

call(z) = where(z > 0, z, alpha * z); fastpath gives the almost-everywhere tower (n = 1 step-with-slope, n >= 2 -> 0).

Parameters

Name	Type & default	Contract
negative_slopepositional-or-keyword	floatDefault: 0.25	See the operation contract and implementation below.
learnablekeyword-only	boolDefault: True	See the operation contract and implementation below.
kwargsvariadic keyword	AnyRequired	See the operation contract and implementation below.

Executed regression · learnable prelu alpha trainable

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_learnable_prelu_alpha_trainable() -> None:
    lp = LearnablePReLU(0.25, learnable=True)
    assert len(lp.trainable_weights) == 1
    assert lp.trainable_weights[0].path.endswith("alpha")
    frozen = LearnablePReLU(0.25, learnable=False)
    assert frozen.trainable_weights == []

python -m pytest "packages/omnibias-keras/tests/test_tempered_blocks.py::test_learnable_prelu_alpha_trainable" -q

Regression example · tempered frozen beta not trainable

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_frozen_beta_not_trainable() -> None:
    frozen = TemperedActivation("softplus", beta=2.0, learnable_beta=False)
    assert frozen.trainable_weights == []

python -m pytest packages/omnibias-keras/tests/test_tempered_blocks.py -k test_tempered_frozen_beta_not_trainable -q

Fields & methods

__init__method

__init__(self, negative_slope: float=0.25, *, learnable: bool=True, **kwargs: Any) -> None

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

Parameters

Name	Type & default	Contract
negative_slopepositional-or-keyword	floatDefault: 0.25	See the operation contract and implementation below.
learnablekeyword-only	boolDefault: True	See the operation contract and implementation below.
kwargsvariadic keyword	AnyRequired	See the operation contract and implementation below.

Returns: None

Executed regression · learnable prelu alpha trainable

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_learnable_prelu_alpha_trainable() -> None:
    lp = LearnablePReLU(0.25, learnable=True)
    assert len(lp.trainable_weights) == 1
    assert lp.trainable_weights[0].path.endswith("alpha")
    frozen = LearnablePReLU(0.25, learnable=False)
    assert frozen.trainable_weights == []

python -m pytest "packages/omnibias-keras/tests/test_tempered_blocks.py::test_learnable_prelu_alpha_trainable" -q

Method implementation

    def __init__(
        self,
        negative_slope: float = 0.25,
        *,
        learnable: bool = True,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        self.init_negative_slope = float(negative_slope)
        self.learnable = learnable
        self.alpha = self.add_weight(
            name="alpha",
            shape=(),
            initializer=initializers.Constant(float(negative_slope)),
            trainable=learnable,
            dtype=self.variable_dtype,
        )
        self.built = True

View source ↗

fastpathmethod

fastpath(self, z: Any, n: int) -> Any

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

Parameters

Name	Type & default	Contract
zpositional-or-keyword	AnyRequired	See the operation contract and implementation below.
npositional-or-keyword	intRequired	See the operation contract and implementation below.

Returns: Any

Validation & failure cases

• When n < 0, raises ValueError(f'order n must be >= 0, got {n}.').

Executed regression · learnable prelu forward and tower

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_learnable_prelu_forward_and_tower() -> None:
    lp = LearnablePReLU(0.25)
    zz = np.array([-2.0, -0.5, 0.5, 2.0])
    z = ops.convert_to_tensor(zz)
    np.testing.assert_allclose(_np(lp(z)), np.where(zz > 0, zz, 0.25 * zz))
    np.testing.assert_allclose(_np(lp.fastpath(z, 1)), np.where(zz > 0, 1.0, 0.25))
    assert np.all(_np(lp.fastpath(z, 2)) == 0.0)

python -m pytest "packages/omnibias-keras/tests/test_tempered_blocks.py::test_learnable_prelu_forward_and_tower" -q

Method implementation

    def fastpath(self, z: Any, n: int) -> Any:
        if n < 0:
            raise ValueError(f"order n must be >= 0, got {n}.")
        if n == 0:
            return ops.where(z > 0, z, self.alpha * z)
        if n == 1:
            return ops.where(z > 0, ops.ones_like(z), self.alpha * ops.ones_like(z))
        return ops.zeros_like(z)

View source ↗

callmethod

call(self, z: Any) -> Any

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

Parameters

Name	Type & default	Contract
zpositional-or-keyword	AnyRequired	See the operation contract and implementation below.

Returns: Any

Executed regression · learnable prelu forward and tower

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_learnable_prelu_forward_and_tower() -> None:
    lp = LearnablePReLU(0.25)
    zz = np.array([-2.0, -0.5, 0.5, 2.0])
    z = ops.convert_to_tensor(zz)
    np.testing.assert_allclose(_np(lp(z)), np.where(zz > 0, zz, 0.25 * zz))
    np.testing.assert_allclose(_np(lp.fastpath(z, 1)), np.where(zz > 0, 1.0, 0.25))
    assert np.all(_np(lp.fastpath(z, 2)) == 0.0)

python -m pytest "packages/omnibias-keras/tests/test_tempered_blocks.py::test_learnable_prelu_forward_and_tower" -q

Method implementation

    def call(self, z: Any) -> Any:
        return self.fastpath(z, 0)

View source ↗

get_configmethod

get_config(self) -> dict[str, Any]

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

Returns: dict[str, Any]

Method implementation

    def get_config(self) -> dict[str, Any]:
        config = super().get_config()
        config.update(
            {
                "negative_slope": self.init_negative_slope,
                "learnable": self.learnable,
            }
        )
        return config

View source ↗

init_negative_slopeattribute

init_negative_slope = float(negative_slope)

Public instance attribute assigned by __init__. The source assignment is init_negative_slope = float(negative_slope); it describes instance state, not a literal constructor default.

View source ↗

learnableattribute

learnable: bool = learnable

Public instance attribute assigned by __init__. The source assignment is learnable = learnable; it describes instance state, not a literal constructor default.

View source ↗

alphaattribute

alpha = self.add_weight(name='alpha', shape=(), initializer=initializers.Constant(float(negative_slope)), trainable=learnable, dtype=self.variable_dtype)

Public instance attribute assigned by __init__. The source assignment is alpha = self.add_weight(name='alpha', shape=(), initializer=initializers.Constant(float(negative_slope)), trainable=learnable, dtype=self.variable_dtype); it describes instance state, not a literal constructor default.

View source ↗

builtattribute

built = True

Public instance attribute assigned by __init__. The source assignment is built = True; it describes instance state, not a literal constructor default.

View source ↗

Read the implementation

class LearnablePReLU(layers.Layer):
    """Leaky ReLU with a (learnable) negative slope ``alpha`` on the a.e. tower.

    ``call(z) = where(z > 0, z, alpha * z)``; ``fastpath`` gives the
    almost-everywhere tower (``n = 1`` step-with-slope, ``n >= 2 -> 0``).
    """

    def __init__(
        self,
        negative_slope: float = 0.25,
        *,
        learnable: bool = True,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        self.init_negative_slope = float(negative_slope)
        self.learnable = learnable
        self.alpha = self.add_weight(
            name="alpha",
            shape=(),
            initializer=initializers.Constant(float(negative_slope)),
            trainable=learnable,
            dtype=self.variable_dtype,
        )
        self.built = True

    def fastpath(self, z: Any, n: int) -> Any:
        if n < 0:
            raise ValueError(f"order n must be >= 0, got {n}.")
        if n == 0:
            return ops.where(z > 0, z, self.alpha * z)
        if n == 1:
            return ops.where(z > 0, ops.ones_like(z), self.alpha * ops.ones_like(z))
        return ops.zeros_like(z)

    def call(self, z: Any) -> Any:
        return self.fastpath(z, 0)

    def get_config(self) -> dict[str, Any]:
        config = super().get_config()
        config.update(
            {
                "negative_slope": self.init_negative_slope,
                "learnable": self.learnable,
            }
        )
        return config

Other public import paths

• omnibias.keras.LearnablePReLU
• omnibias.keras.tempered_blocks.LearnablePReLU

CLASS

TemperedActivation

omnibias.keras.TemperedActivationSource ↗

class TemperedActivation(layers.Layer)

Beta-tempered smooth surrogate of a base activation's tower.

Parameters

base : str or ActivationSpec, default "softplus" Base activation carrying a fastpath. The surrogate tower is beta**(n - p) * base.fastpath(beta * z, n). beta : float, default 1.0 Initial temperature (larger = sharper; beta -> inf approaches the hard activation). scale : {"one_over_beta", "unit"}, default "one_over_beta" "one_over_beta" (p = 1) for softplus(beta z)/beta -> relu; "unit" (p = 0) for bounded surrogates like sigmoid / tanh. learnable_beta : bool, default False If True, beta is a trainable weight; otherwise it is frozen.

Parameters

Name	Type & default	Contract
basepositional-or-keyword	str | ActivationSpec[Any]Default: 'softplus'	str or ActivationSpec, default "softplus" Base activation carrying a fastpath. The surrogate tower is beta**(n - p) * base.fastpath(beta * z, n).
betapositional-or-keyword	floatDefault: 1.0	float, default 1.0 Initial temperature (larger = sharper; beta -> inf approaches the hard activation).
scalekeyword-only	strDefault: 'one_over_beta'	{"one_over_beta", "unit"}, default "one_over_beta" "one_over_beta" (p = 1) for softplus(beta z)/beta -> relu; "unit" (p = 0) for bounded surrogates like sigmoid / tanh.
learnable_betakeyword-only	boolDefault: False	bool, default False If True, beta is a trainable weight; otherwise it is frozen.
kwargsvariadic keyword	AnyRequired	See the operation contract and implementation below.

Executed regression · tempered rejects bad scale

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_tempered_rejects_bad_scale() -> None:
    with pytest.raises(ValueError):
        TemperedActivation("softplus", beta=1.0, scale="bogus")

python -m pytest "packages/omnibias-keras/tests/test_tempered_blocks.py::test_tempered_rejects_bad_scale" -q

Regression example · tempered frozen beta not trainable

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_frozen_beta_not_trainable() -> None:
    frozen = TemperedActivation("softplus", beta=2.0, learnable_beta=False)
    assert frozen.trainable_weights == []

python -m pytest packages/omnibias-keras/tests/test_tempered_blocks.py -k test_tempered_frozen_beta_not_trainable -q

Fields & methods

__init__method

__init__(self, base: str | ActivationSpec[Any]='softplus', beta: float=1.0, *, scale: str='one_over_beta', learnable_beta: bool=False, **kwargs: Any) -> None

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

Parameters

Name	Type & default	Contract
basepositional-or-keyword	str | ActivationSpec[Any]Default: 'softplus'	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.
scalekeyword-only	strDefault: 'one_over_beta'	See the operation contract and implementation below.
learnable_betakeyword-only	boolDefault: False	See the operation contract and implementation below.
kwargsvariadic keyword	AnyRequired	See the operation contract and implementation below.

Returns: None

Validation & failure cases

• When spec.fastpath is None, raises ValueError(f'TemperedActivation requires a base with a fastpath; {spec.name!r} has none.').

Executed regression · tempered rejects bad scale

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_tempered_rejects_bad_scale() -> None:
    with pytest.raises(ValueError):
        TemperedActivation("softplus", beta=1.0, scale="bogus")

python -m pytest "packages/omnibias-keras/tests/test_tempered_blocks.py::test_tempered_rejects_bad_scale" -q

Method implementation

    def __init__(
        self,
        base: str | ActivationSpec[Any] = "softplus",
        beta: float = 1.0,
        *,
        scale: str = "one_over_beta",
        learnable_beta: bool = False,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        spec = base if isinstance(base, ActivationSpec) else get_activation(base)
        if spec.fastpath is None:
            raise ValueError(
                f"TemperedActivation requires a base with a fastpath; {spec.name!r} has none."
            )
        if scale == "one_over_beta":
            self._scale_power = 1
        elif scale == "unit":
            self._scale_power = 0
        else:
            raise ValueError(f"scale must be 'unit' or 'one_over_beta', got {scale!r}")
        self._base_name = base if isinstance(base, str) else spec.name
        self.scale = scale
        self.init_beta = float(beta)
        self.learnable_beta = learnable_beta
        self._base_fastpath = spec.fastpath
        self.beta = self.add_weight(
            name="beta",
            shape=(),
            initializer=initializers.Constant(float(beta)),
            trainable=learnable_beta,
            dtype=self.variable_dtype,
        )
        self.built = True

View source ↗

fastpathmethod

fastpath(self, z: Any, n: int) -> Any

Closed-form n-th derivative of the tempered surrogate at current beta.

Parameters

Name	Type & default	Contract
zpositional-or-keyword	AnyRequired	See the operation contract and implementation below.
npositional-or-keyword	intRequired	See the operation contract and implementation below.

Returns: Any

Executed regression · tempered call is forward

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_tempered_call_is_forward() -> None:
    mod = TemperedActivation("sigmoid", beta=1.5, scale="unit")
    z = _z()
    np.testing.assert_allclose(_np(mod(z)), _np(mod.fastpath(z, 0)), rtol=1e-6, atol=1e-6)

python -m pytest "packages/omnibias-keras/tests/test_tempered_blocks.py::test_tempered_call_is_forward" -q

Method implementation

    def fastpath(self, z: Any, n: int) -> Any:
        """Closed-form ``n``-th derivative of the tempered surrogate at current ``beta``."""
        kernel = make_tempered_fastpath(
            self._base_fastpath, self.beta, scale_power=self._scale_power
        )
        return kernel(z, n)

View source ↗

callmethod

call(self, z: Any) -> Any

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

Parameters

Name	Type & default	Contract
zpositional-or-keyword	AnyRequired	See the operation contract and implementation below.

Returns: Any

Executed regression · tempered call is forward

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_tempered_call_is_forward() -> None:
    mod = TemperedActivation("sigmoid", beta=1.5, scale="unit")
    z = _z()
    np.testing.assert_allclose(_np(mod(z)), _np(mod.fastpath(z, 0)), rtol=1e-6, atol=1e-6)

python -m pytest "packages/omnibias-keras/tests/test_tempered_blocks.py::test_tempered_call_is_forward" -q

Method implementation

    def call(self, z: Any) -> Any:
        return self.fastpath(z, 0)

View source ↗

get_configmethod

get_config(self) -> dict[str, Any]

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

Returns: dict[str, Any]

Executed regression · tempered config round trip

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_tempered_config_round_trip() -> None:
    mod = TemperedActivation("sigmoid", beta=2.5, scale="unit", learnable_beta=True)
    clone = TemperedActivation.from_config(mod.get_config())
    z = _z()
    np.testing.assert_allclose(_np(clone(z)), _np(mod(z)), rtol=1e-6, atol=1e-6)
    assert clone.scale == "unit"

python -m pytest "packages/omnibias-keras/tests/test_tempered_blocks.py::test_tempered_config_round_trip" -q

Method implementation

    def get_config(self) -> dict[str, Any]:
        config = super().get_config()
        config.update(
            {
                "base": self._base_name,
                "beta": self.init_beta,
                "scale": self.scale,
                "learnable_beta": self.learnable_beta,
            }
        )
        return config

View source ↗

scaleattribute

scale: str = scale

"one_over_beta" (p = 1) for softplus(beta z)/beta -> relu; "unit" (p = 0) for bounded surrogates like sigmoid / tanh.

View source ↗

init_betaattribute

init_beta = float(beta)

Public instance attribute assigned by __init__. The source assignment is init_beta = float(beta); it describes instance state, not a literal constructor default.

View source ↗

learnable_betaattribute

learnable_beta: bool = learnable_beta

If True, beta is a trainable weight; otherwise it is frozen.

View source ↗

betaattribute

beta = self.add_weight(name='beta', shape=(), initializer=initializers.Constant(float(beta)), trainable=learnable_beta, dtype=self.variable_dtype)

Initial temperature (larger = sharper; beta -> inf approaches the hard activation).

View source ↗

builtattribute

built = True

Public instance attribute assigned by __init__. The source assignment is built = True; it describes instance state, not a literal constructor default.

View source ↗

Read the implementation

class TemperedActivation(layers.Layer):
    """Beta-tempered smooth surrogate of a base activation's tower.

    Parameters
    ----------
    base : str or :class:`ActivationSpec`, default ``"softplus"``
        Base activation carrying a fastpath. The surrogate tower is
        ``beta**(n - p) * base.fastpath(beta * z, n)``.
    beta : float, default 1.0
        Initial temperature (larger = sharper; ``beta -> inf`` approaches the
        hard activation).
    scale : {"one_over_beta", "unit"}, default ``"one_over_beta"``
        ``"one_over_beta"`` (``p = 1``) for ``softplus(beta z)/beta -> relu``;
        ``"unit"`` (``p = 0``) for bounded surrogates like ``sigmoid`` / ``tanh``.
    learnable_beta : bool, default False
        If True, ``beta`` is a trainable weight; otherwise it is frozen.
    """

    def __init__(
        self,
        base: str | ActivationSpec[Any] = "softplus",
        beta: float = 1.0,
        *,
        scale: str = "one_over_beta",
        learnable_beta: bool = False,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        spec = base if isinstance(base, ActivationSpec) else get_activation(base)
        if spec.fastpath is None:
            raise ValueError(
                f"TemperedActivation requires a base with a fastpath; {spec.name!r} has none."
            )
        if scale == "one_over_beta":
            self._scale_power = 1
        elif scale == "unit":
            self._scale_power = 0
        else:
            raise ValueError(f"scale must be 'unit' or 'one_over_beta', got {scale!r}")
        self._base_name = base if isinstance(base, str) else spec.name
        self.scale = scale
        self.init_beta = float(beta)
        self.learnable_beta = learnable_beta
        self._base_fastpath = spec.fastpath
        self.beta = self.add_weight(
            name="beta",
            shape=(),
            initializer=initializers.Constant(float(beta)),
            trainable=learnable_beta,
            dtype=self.variable_dtype,
        )
        self.built = True

    def fastpath(self, z: Any, n: int) -> Any:
        """Closed-form ``n``-th derivative of the tempered surrogate at current ``beta``."""
        kernel = make_tempered_fastpath(
            self._base_fastpath, self.beta, scale_power=self._scale_power
        )
        return kernel(z, n)

    def call(self, z: Any) -> Any:
        return self.fastpath(z, 0)

    def get_config(self) -> dict[str, Any]:
        config = super().get_config()
        config.update(
            {
                "base": self._base_name,
                "beta": self.init_beta,
                "scale": self.scale,
                "learnable_beta": self.learnable_beta,
            }
        )
        return config

Related APIsomnibias.core.make_tempered_fastpath↗omnibias.keras.get_activation↗

Other public import paths

• omnibias.keras.TemperedActivation
• omnibias.keras.tempered_blocks.TemperedActivation

← Package overviewRunnable learning examples →