Source code for embedl_deploy._internal.lattice.modules.activation

# Copyright (C) 2026 Embedl AB

"""Lattice-conforming activation modules.

Lattice hardware supports ``ReLU`` and ``LeakyReLU`` as the activation
in a CBSR block.  The ``LeakyReLU`` negative slope is fixed at
``1/16`` (0.0625).
:class:`~embedl_deploy._internal.lattice.modules.activation.LatticeLeakyReLU`
subclasses :class:`~torch.nn.LeakyReLU` and snaps the slope to that value.
"""

from typing import TypeAlias

from torch import nn

from embedl_deploy._internal.core.modules import ConvertedModule

#: The single negative-slope value accepted by Lattice hardware.
LATTICE_LEAKY_RELU_SLOPE: float = 0.0625




class LatticeLeakyReLU(ConvertedModule, nn.LeakyReLU):
    """``LeakyReLU`` snapped to Lattice's supported slope.

    Lattice hardware implements leaky ReLU with a fixed negative slope
    of ``1/16`` (0.0625).  The constructor always sets the slope to the
    supported value; the conversion pattern logs a warning when the
    source slope differs.

    Subclasses
    :class:`~embedl_deploy._internal.core.modules.ConvertedModule` so
    that FX tracing keeps this module as a leaf node.
    """

    #: The single permitted negative slope.
    NEGATIVE_SLOPE: float = LATTICE_LEAKY_RELU_SLOPE



    @classmethod
    def is_compatible(cls, act: nn.LeakyReLU) -> bool:
        """Return ``True`` when `act` already uses the supported slope.

        :param act:
            Activation to check.
        :returns:
            ``True`` when the negative slope of `act` equals
            ``NEGATIVE_SLOPE``; ``False`` otherwise.
        """
        return act.negative_slope == cls.NEGATIVE_SLOPE


    def __init__(self, act: nn.LeakyReLU) -> None:
        """Create a ``LatticeLeakyReLU`` from an arbitrary ``LeakyReLU``.

        :param act:
            Source activation.  Its ``inplace`` flag is preserved;
            the negative slope is replaced with the hardware-fixed
            value.
        """
        super().__init__(
            negative_slope=self.NEGATIVE_SLOPE,
            inplace=act.inplace,
        )



#: Activation types accepted by Lattice CBSR blocks.
LatticeActivationLike: TypeAlias = nn.ReLU | LatticeLeakyReLU