Source code for embedl_deploy._internal.lattice.patterns.conversions.activation

# Copyright (C) 2026 Embedl AB

"""Activation conversion patterns for Lattice."""

import logging

from torch import fx, nn

from embedl_deploy._internal.core.patterns.main import Pattern, Phase
from embedl_deploy._internal.core.tree.types import Graft, Tree, TreeMatch
from embedl_deploy._internal.core.tree.utils import get_module, resolve_module
from embedl_deploy._internal.lattice.modules.activation import (
    LatticeLeakyReLU,
)

_LOG = logging.getLogger(__name__)


def _is_nonconforming_leaky_relu(node: fx.Node) -> bool:
    """Return ``True`` when `node` is a ``LeakyReLU`` that needs rewriting."""
    mod = get_module(node)
    if not isinstance(mod, nn.LeakyReLU):
        return False
    if isinstance(mod, LatticeLeakyReLU):
        return False
    return not LatticeLeakyReLU.is_compatible(mod)


def _make_lattice_leaky_relu(
    tree_match: TreeMatch,
) -> tuple[nn.Module, ...]:
    """Build the replacement ``LeakyReLU`` for the matched node."""
    node = tree_match.get_node(0)
    act = resolve_module(node, nn.LeakyReLU)
    new_act = LatticeLeakyReLU(act)
    if act.negative_slope != new_act.negative_slope:
        _LOG.warning(
            "%s: LeakyReLU negative_slope snapped from %s to %s — "
            "output will differ, retraining is recommended",
            node.name,
            act.negative_slope,
            new_act.negative_slope,
        )
    return (new_act,)




class LatticeLeakyReLUPattern(Pattern):
    """Snap a ``LeakyReLU`` negative slope to Lattice's supported value.

    Lattice hardware implements leaky ReLU with a fixed negative slope of
    ``1/16`` (0.0625). Any other slope is replaced by that value and a warning
    is logged advising retraining.
    """

    is_numerically_equivalent = False
    phase = Phase.CONVERSION
    tree: Tree = (_is_nonconforming_leaky_relu,)
    graft: Graft = (_make_lattice_leaky_relu,)



def _make_relu_from_activation(
    tree_match: TreeMatch,
) -> tuple[nn.Module, ...]:
    """Build a ``ReLU`` replacement for the matched activation."""
    node = tree_match.get_node(0)
    act = resolve_module(node, nn.Module)
    _LOG.warning(
        "%s: %s replaced with ReLU — output will differ, retraining is recommended",
        node.name,
        type(act).__name__,
    )
    return (nn.ReLU(inplace=getattr(act, "inplace", False)),)




class LatticeSiLUToReLUPattern(Pattern):
    """Replace ``SiLU`` with ``ReLU``.

    Lattice hardware does not support ``SiLU`` activations. Any ``SiLU``
    instance is replaced by a ``ReLU`` and a warning is logged advising
    retraining.
    """

    is_numerically_equivalent = False
    phase = Phase.CONVERSION
    tree: Tree = (lambda node: isinstance(get_module(node), nn.SiLU),)
    graft: Graft = (_make_relu_from_activation,)





class LatticeGELUToReLUPattern(Pattern):
    """Replace ``GELU`` with ``ReLU``.

    Lattice hardware does not support ``GELU`` activations. Any ``GELU``
    instance is replaced by a ``ReLU`` and a warning is logged advising
    retraining.
    """

    is_numerically_equivalent = False
    phase = Phase.CONVERSION
    tree: Tree = (lambda node: isinstance(get_module(node), nn.GELU),)
    graft: Graft = (_make_relu_from_activation,)





class LatticeReLU6ToReLUPattern(Pattern):
    """Replace ``ReLU6`` with ``ReLU``.

    Lattice hardware does not support ``ReLU6`` (Clip) activations. Any
    ``ReLU6`` instance is replaced by a ``ReLU`` and a warning is logged
    advising retraining.
    """

    is_numerically_equivalent = False
    phase = Phase.CONVERSION
    tree: Tree = (lambda node: isinstance(get_module(node), nn.ReLU6),)
    graft: Graft = (_make_relu_from_activation,)