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

# Copyright (C) 2026 Embedl AB

"""Fused pointwise element-wise operation modules for Lattice."""

from typing import cast

import torch

from embedl_deploy._internal.core.modules import FusedModule
from embedl_deploy._internal.core.quantize.config import Precision
from embedl_deploy._internal.lattice.modules.activation import (
    LatticeActivationLike,
)
from embedl_deploy._internal.lattice.modules.quant import LatticeQuant




class LatticeFusedAddAct(FusedModule):
    """Fused ``add(·, residual) -> [Activation]`` for Lattice hardware.

    Fuses the element-wise addition of two branches with an optional trailing
    activation (``ReLU`` or ``LatticeLeakyReLU``) into a single block. This
    captures the residual merge point in ResNet-style architectures where the
    convolution path has already been fused separately.

    Unlike :class:`~embedl_deploy._internal.lattice.modules.conv.LatticeCBSR`
    there is no convolution, batch normalization, or weight tensor, so no
    ``WeightFakeQuantize`` is attached.

    :param act:
        Optional ``ReLU`` or ``LatticeLeakyReLU`` activation applied after
        the element-wise addition.
    :param output_precision:
        Precision for the output quantizer.
    :param output_fixed_calibration:
        Optional fixed ``(scale, zero_point)`` pair for the output
        quantizer.
    """

    inputs_to_quantize: set[int] = set()

    def __init__(
        self,
        act: LatticeActivationLike | None = None,
        *,
        output_precision: Precision = Precision.INT8,
        output_fixed_calibration: tuple[float, int] | None = None,
    ) -> None:
        super().__init__()
        self.act = act
        self.output_quant_stub = LatticeQuant(
            {self},
            precision=output_precision,
            fixed_calibration=output_fixed_calibration,
        )



    def forward(self, x: torch.Tensor, residual: torch.Tensor) -> torch.Tensor:
        """Apply ``add(x, residual) -> [act]``."""
        x = x + residual
        if self.act is not None:
            x = cast(torch.Tensor, self.act(x))
        return x


    def __repr__(self) -> str:
        act_info = ""
        if self.act is not None:
            act_info = f", act={type(self.act).__name__}"
        return f"LatticeFusedAddAct(add{act_info})"