API REFERENCE/omnibias-torch

omnibias.torch.tempered_blocks

0.5.0rc1Release candidatePython >=3.10Apache-2.0

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

import omnibias.torch.tempered_blocks

Module contract & mathematical conventions

Learnable-temperature activation modules.

Thin nn.Module wrappers that hold a differentiable temperature beta (or negative slope alpha) and call the closed-form tempered / piecewise tower each forward -- mirroring the learnable-beta pattern in omnibias-binary's binarize. The whole derivative tower stays available through fastpath(z, n), so these compose with the operator layers.

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

beta (and alpha) should stay positive; wrap with your own reparameterisation (softplus / exp) 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.torch.LearnablePReLUSource ↗

class LearnablePReLU(nn.Module)

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

forward(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.

Executed regression · learnable prelu

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() -> None:
    z = torch.tensor([-2.0, -0.5, 0.5, 2.0], dtype=torch.float64)
    lp = LearnablePReLU(0.25, learnable=True)
    assert isinstance(lp.alpha, torch.nn.Parameter)
    torch.testing.assert_close(lp(z), torch.where(z > 0, z, 0.25 * z))
    torch.testing.assert_close(
        lp.fastpath(z, 1), torch.where(z > 0, torch.ones_like(z), torch.full_like(z, 0.25))
    )
    assert torch.equal(lp.fastpath(z, 2), torch.zeros_like(z))
    lp(z).sum().backward()
    assert lp.alpha.grad is not None

python -m pytest "packages/omnibias-torch/tests/test_tempered.py::test_learnable_prelu" -q

Fields & methods

alphafield

alpha: Tensor

Declared Tensor field on LearnablePReLU. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

Regression example · learnable prelu

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() -> None:
    z = torch.tensor([-2.0, -0.5, 0.5, 2.0], dtype=torch.float64)
    lp = LearnablePReLU(0.25, learnable=True)
    assert isinstance(lp.alpha, torch.nn.Parameter)
    torch.testing.assert_close(lp(z), torch.where(z > 0, z, 0.25 * z))
    torch.testing.assert_close(
        lp.fastpath(z, 1), torch.where(z > 0, torch.ones_like(z), torch.full_like(z, 0.25))
    )
    assert torch.equal(lp.fastpath(z, 2), torch.zeros_like(z))
    lp(z).sum().backward()
    assert lp.alpha.grad is not None

python -m pytest packages/omnibias-torch/tests/test_tempered.py -k test_learnable_prelu -q

View source ↗

__init__method

__init__(self, negative_slope: float=0.25, *, learnable: bool=True) -> 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.

Returns: None

Executed regression · learnable prelu

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() -> None:
    z = torch.tensor([-2.0, -0.5, 0.5, 2.0], dtype=torch.float64)
    lp = LearnablePReLU(0.25, learnable=True)
    assert isinstance(lp.alpha, torch.nn.Parameter)
    torch.testing.assert_close(lp(z), torch.where(z > 0, z, 0.25 * z))
    torch.testing.assert_close(
        lp.fastpath(z, 1), torch.where(z > 0, torch.ones_like(z), torch.full_like(z, 0.25))
    )
    assert torch.equal(lp.fastpath(z, 2), torch.zeros_like(z))
    lp(z).sum().backward()
    assert lp.alpha.grad is not None

python -m pytest "packages/omnibias-torch/tests/test_tempered.py::test_learnable_prelu" -q

Method implementation

    def __init__(self, negative_slope: float = 0.25, *, learnable: bool = True) -> None:
        super().__init__()
        alpha_init = torch.as_tensor(float(negative_slope))
        if learnable:
            self.alpha = nn.Parameter(alpha_init)
        else:
            self.register_buffer("alpha", alpha_init)

View source ↗

fastpathmethod

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

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

Parameters

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

Returns: Tensor

Validation & failure cases

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

Executed regression · learnable prelu

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() -> None:
    z = torch.tensor([-2.0, -0.5, 0.5, 2.0], dtype=torch.float64)
    lp = LearnablePReLU(0.25, learnable=True)
    assert isinstance(lp.alpha, torch.nn.Parameter)
    torch.testing.assert_close(lp(z), torch.where(z > 0, z, 0.25 * z))
    torch.testing.assert_close(
        lp.fastpath(z, 1), torch.where(z > 0, torch.ones_like(z), torch.full_like(z, 0.25))
    )
    assert torch.equal(lp.fastpath(z, 2), torch.zeros_like(z))
    lp(z).sum().backward()
    assert lp.alpha.grad is not None

python -m pytest "packages/omnibias-torch/tests/test_tempered.py::test_learnable_prelu" -q

Method implementation

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

View source ↗

forwardmethod

forward(self, z: Tensor) -> Tensor

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

Parameters

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

Returns: Tensor

Executed regression · learnable prelu

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() -> None:
    z = torch.tensor([-2.0, -0.5, 0.5, 2.0], dtype=torch.float64)
    lp = LearnablePReLU(0.25, learnable=True)
    assert isinstance(lp.alpha, torch.nn.Parameter)
    torch.testing.assert_close(lp(z), torch.where(z > 0, z, 0.25 * z))
    torch.testing.assert_close(
        lp.fastpath(z, 1), torch.where(z > 0, torch.ones_like(z), torch.full_like(z, 0.25))
    )
    assert torch.equal(lp.fastpath(z, 2), torch.zeros_like(z))
    lp(z).sum().backward()
    assert lp.alpha.grad is not None

python -m pytest "packages/omnibias-torch/tests/test_tempered.py::test_learnable_prelu" -q

Method implementation

    def forward(self, z: Tensor) -> Tensor:
        return self.fastpath(z, 0)

View source ↗

extra_reprmethod

extra_repr(self) -> str

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

Returns: str

Method implementation

    def extra_repr(self) -> str:
        return f"learnable={isinstance(self.alpha, nn.Parameter)}"

View source ↗

Read the implementation

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

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

    alpha: Tensor

    def __init__(self, negative_slope: float = 0.25, *, learnable: bool = True) -> None:
        super().__init__()
        alpha_init = torch.as_tensor(float(negative_slope))
        if learnable:
            self.alpha = nn.Parameter(alpha_init)
        else:
            self.register_buffer("alpha", alpha_init)

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

    def forward(self, z: Tensor) -> Tensor:
        return self.fastpath(z, 0)

    def extra_repr(self) -> str:
        return f"learnable={isinstance(self.alpha, nn.Parameter)}"

Other public import paths

• omnibias.torch.LearnablePReLU
• omnibias.torch.tempered_blocks.LearnablePReLU

CLASS

TemperedActivation

omnibias.torch.TemperedActivationSource ↗

class TemperedActivation(nn.Module)

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

Parameters

base : str or ActivationSpec, default "softplus" Base activation carrying a fastpath (e.g. "softplus", "sigmoid", "tanh"). 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 an nn.Parameter (differentiable temperature); otherwise it is a frozen buffer.

Parameters

Name	Type & default	Contract
basepositional-or-keyword	str | ActivationSpec[Tensor]Default: 'softplus'	str or ActivationSpec, default "softplus" Base activation carrying a fastpath (e.g. "softplus", "sigmoid", "tanh"). 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 an nn.Parameter (differentiable temperature); otherwise it is a frozen buffer.

Executed regression · frozen beta is buffer not parameter

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_frozen_beta_is_buffer_not_parameter() -> None:
    mod = TemperedActivation("sigmoid", beta=3.0, scale="unit", learnable_beta=False)
    assert not isinstance(mod.beta, torch.nn.Parameter)
    assert "beta" in dict(mod.named_buffers())

python -m pytest "packages/omnibias-torch/tests/test_tempered.py::test_frozen_beta_is_buffer_not_parameter" -q

Regression example · tempered activation module matches spec

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_module_matches_spec() -> None:
    z = torch.linspace(-3.0, 3.0, 41, dtype=torch.float64)
    mod = TemperedActivation("softplus", beta=2.0, scale="one_over_beta")
    spec = make_soft_relu_spec(2.0)
    for n in range(0, 5):
        torch.testing.assert_close(mod.fastpath(z, n), spec.fastpath(z, n))
    torch.testing.assert_close(mod(z), spec.forward(z))

python -m pytest packages/omnibias-torch/tests/test_tempered.py -k test_tempered_activation_module_matches_spec -q

Fields & methods

betafield

beta: Tensor

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

Regression example · frozen beta is buffer not parameter

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_frozen_beta_is_buffer_not_parameter() -> None:
    mod = TemperedActivation("sigmoid", beta=3.0, scale="unit", learnable_beta=False)
    assert not isinstance(mod.beta, torch.nn.Parameter)
    assert "beta" in dict(mod.named_buffers())

python -m pytest packages/omnibias-torch/tests/test_tempered.py -k test_frozen_beta_is_buffer_not_parameter -q

View source ↗

__init__method

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

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

Parameters

Name	Type & default	Contract
basepositional-or-keyword	str | ActivationSpec[Tensor]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.

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 · frozen beta is buffer not parameter

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_frozen_beta_is_buffer_not_parameter() -> None:
    mod = TemperedActivation("sigmoid", beta=3.0, scale="unit", learnable_beta=False)
    assert not isinstance(mod.beta, torch.nn.Parameter)
    assert "beta" in dict(mod.named_buffers())

python -m pytest "packages/omnibias-torch/tests/test_tempered.py::test_frozen_beta_is_buffer_not_parameter" -q

Method implementation

    def __init__(
        self,
        base: str | ActivationSpec[Tensor] = "softplus",
        beta: float = 1.0,
        *,
        scale: str = "one_over_beta",
        learnable_beta: bool = False,
    ) -> None:
        super().__init__()
        spec = 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 = spec.name
        self.scale = scale
        self._base_fastpath = spec.fastpath
        beta_init = torch.as_tensor(float(beta))
        if learnable_beta:
            self.beta = nn.Parameter(beta_init)
        else:
            self.register_buffer("beta", beta_init)

View source ↗

fastpathmethod

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

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

Parameters

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

Returns: Tensor

Executed regression · tempered activation module matches spec

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_activation_module_matches_spec() -> None:
    z = torch.linspace(-3.0, 3.0, 41, dtype=torch.float64)
    mod = TemperedActivation("softplus", beta=2.0, scale="one_over_beta")
    spec = make_soft_relu_spec(2.0)
    for n in range(0, 5):
        torch.testing.assert_close(mod.fastpath(z, n), spec.fastpath(z, n))
    torch.testing.assert_close(mod(z), spec.forward(z))

python -m pytest "packages/omnibias-torch/tests/test_tempered.py::test_tempered_activation_module_matches_spec" -q

Method implementation

    def fastpath(self, z: Tensor, n: int) -> Tensor:
        """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 ↗

forwardmethod

forward(self, z: Tensor) -> Tensor

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

Parameters

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

Returns: Tensor

Executed regression · tempered activation module matches spec

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_activation_module_matches_spec() -> None:
    z = torch.linspace(-3.0, 3.0, 41, dtype=torch.float64)
    mod = TemperedActivation("softplus", beta=2.0, scale="one_over_beta")
    spec = make_soft_relu_spec(2.0)
    for n in range(0, 5):
        torch.testing.assert_close(mod.fastpath(z, n), spec.fastpath(z, n))
    torch.testing.assert_close(mod(z), spec.forward(z))

python -m pytest "packages/omnibias-torch/tests/test_tempered.py::test_tempered_activation_module_matches_spec" -q

Method implementation

    def forward(self, z: Tensor) -> Tensor:
        return self.fastpath(z, 0)

View source ↗

extra_reprmethod

extra_repr(self) -> str

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

Returns: str

Method implementation

    def extra_repr(self) -> str:
        learnable = isinstance(self.beta, nn.Parameter)
        return f"base={self.base_name!r}, scale={self.scale!r}, learnable_beta={learnable}"

View source ↗

base_nameattribute

base_name = spec.name

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

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 ↗

Read the implementation

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

    Parameters
    ----------
    base : str or :class:`ActivationSpec`, default ``"softplus"``
        Base activation carrying a fastpath (e.g. ``"softplus"``, ``"sigmoid"``,
        ``"tanh"``). 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 an ``nn.Parameter`` (differentiable temperature);
        otherwise it is a frozen buffer.
    """

    beta: Tensor
    _base_fastpath: NthDerivativeFn

    def __init__(
        self,
        base: str | ActivationSpec[Tensor] = "softplus",
        beta: float = 1.0,
        *,
        scale: str = "one_over_beta",
        learnable_beta: bool = False,
    ) -> None:
        super().__init__()
        spec = 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 = spec.name
        self.scale = scale
        self._base_fastpath = spec.fastpath
        beta_init = torch.as_tensor(float(beta))
        if learnable_beta:
            self.beta = nn.Parameter(beta_init)
        else:
            self.register_buffer("beta", beta_init)

    def fastpath(self, z: Tensor, n: int) -> Tensor:
        """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 forward(self, z: Tensor) -> Tensor:
        return self.fastpath(z, 0)

    def extra_repr(self) -> str:
        learnable = isinstance(self.beta, nn.Parameter)
        return f"base={self.base_name!r}, scale={self.scale!r}, learnable_beta={learnable}"

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

Other public import paths

• omnibias.torch.TemperedActivation
• omnibias.torch.tempered_blocks.TemperedActivation

← Package overviewRunnable learning examples →