# (C) Copyright 2024 Anemoi contributors.
#
# This software is licensed under the terms of the Apache Licence Version 2.0
# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
#
# In applying this licence, ECMWF does not waive the privileges and immunities
# granted to it by virtue of its status as an intergovernmental organisation
# nor does it submit to any jurisdiction.


import logging
import os
from abc import ABC
from abc import abstractmethod
from typing import Optional

import torch
from torch import Tensor
from torch import nn
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import offload_wrapper
from torch.distributed.distributed_c10d import ProcessGroup
from torch_geometric.typing import Adj
from torch_geometric.typing import PairTensor

from anemoi.models.distributed.graph import gather_tensor
from anemoi.models.distributed.graph import shard_tensor
from anemoi.models.distributed.graph import sync_tensor
from anemoi.models.distributed.khop_edges import bipartite_subgraph
from anemoi.models.distributed.khop_edges import drop_unconnected_src_nodes
from anemoi.models.distributed.khop_edges import shard_edges_1hop
from anemoi.models.distributed.shapes import change_channels_in_shape
from anemoi.models.distributed.shapes import get_shard_shapes
from anemoi.models.layers.block import GraphConvMapperBlock
from anemoi.models.layers.block import GraphTransformerMapperBlock
from anemoi.models.layers.block import TransformerMapperBlock
from anemoi.models.layers.mlp import MLP
from anemoi.models.layers.utils import load_layer_kernels
from anemoi.models.layers.utils import maybe_checkpoint
from anemoi.utils.config import DotDict

LOGGER = logging.getLogger(__name__)

# Number of chunks used in inference (https://github.com/ecmwf/anemoi-core/pull/406)
NUM_CHUNKS_INFERENCE = int(os.environ.get("ANEMOI_INFERENCE_NUM_CHUNKS", "1"))
NUM_CHUNKS_INFERENCE_MAPPER = int(os.environ.get("ANEMOI_INFERENCE_NUM_CHUNKS_MAPPER", NUM_CHUNKS_INFERENCE))


class BaseMapper(nn.Module, ABC):
    """Base Mapper from source dimension to destination dimension.

    Subclasses must implement pre_process() and post_process() methods
    specialized for their mapper type.
    """

    def __init__(
        self,
        *,
        in_channels_src: int,
        in_channels_dst: int,
        hidden_dim: int,
        out_channels_dst: Optional[int] = None,
        cpu_offload: bool = False,
        gradient_checkpointing: bool = True,
        layer_kernels: DotDict,
        **kwargs,
    ) -> None:
        """Initialize BaseMapper."""
        super().__init__()

        self.in_channels_src = in_channels_src
        self.in_channels_dst = in_channels_dst
        self.hidden_dim = hidden_dim
        self.out_channels_dst = out_channels_dst
        self.gradient_checkpointing = gradient_checkpointing
        self.layer_factory = load_layer_kernels(layer_kernels)
        self.activation = self.layer_factory.Activation()

        self.proc = NotImplemented

        self.offload_layers(cpu_offload)

    def offload_layers(self, cpu_offload):
        if cpu_offload:
            self.proc = nn.ModuleList([offload_wrapper(x) for x in self.proc])

    @abstractmethod
    def pre_process(self, x, shard_shapes, model_comm_group=None, x_src_is_sharded=False, x_dst_is_sharded=False):
        pass

    @abstractmethod
    def post_process(self, x_dst, shapes_dst=None, model_comm_group=None, keep_x_dst_sharded=False, **kwargs):
        pass

    @abstractmethod
    def forward(
        self,
        x: PairTensor,
        batch_size: int,
        shard_shapes: tuple[tuple[int], tuple[int]],
        edge_attr: Optional[Tensor] = None,
        edge_index: Optional[Adj] = None,
        model_comm_group: Optional[ProcessGroup] = None,
        x_src_is_sharded: bool = False,
        x_dst_is_sharded: bool = False,
        keep_x_dst_sharded: bool = False,
        edge_shard_shapes: Optional[tuple] = None,
        **kwargs,
    ):
        """Forward pass of the mapper.

        Parameters
        ----------
        x : PairTensor
            Input tensor pair (source, destination)
        batch_size : int
            Batch size
        shard_shapes : tuple[tuple[int], tuple[int]]
            Shard shapes for source and destination
        edge_attr : Tensor, optional
            Edge attributes (required for graph-based mappers)
        edge_index : Adj, optional
            Edge indices (required for graph-based mappers)
        model_comm_group : ProcessGroup, optional
            Model communication group
        x_src_is_sharded : bool, optional
            Whether source is already sharded, by default False
        x_dst_is_sharded : bool, optional
            Whether destination is already sharded, by default False
        keep_x_dst_sharded : bool, optional
            Whether to keep destination sharded, by default False
        edge_shard_shapes : tuple, optional
            Shard shapes for edges (shapes_edge_attr, shapes_edge_idx)
            If provided, edges are gathered before use.
        """
        pass


class GraphTransformerBaseMapper(BaseMapper, ABC):
    """Graph Transformer Base Mapper from hidden -> data or data -> hidden."""

    def __init__(
        self,
        *,
        in_channels_src: int,
        in_channels_dst: int,
        hidden_dim: int,
        out_channels_dst: Optional[int] = None,
        num_chunks: int,
        num_heads: int,
        mlp_hidden_ratio: int,
        edge_dim: int,
        qk_norm: bool = False,
        cpu_offload: bool = False,
        layer_kernels: DotDict = None,
        shard_strategy: str = "edges",
        graph_attention_backend: str = "triton",
        edge_pre_mlp: bool = False,
        **kwargs,
    ) -> None:
        """Initialize GraphTransformerBaseMapper.

        Parameters
        ----------
        in_channels_src : int
            Input channels of the source node
        in_channels_dst : int
            Input channels of the destination node
        hidden_dim : int
            Hidden dimension
        out_channels_dst : int, optional
            Output channels of the destination node, by default None
        num_chunks : int
            Number of chunks to split into
        num_heads: int
            Number of heads in transformer
        mlp_hidden_ratio: int
            ratio of mlp hidden dimension to embedding dimension
        edge_dim : int
            Edge feature dimension
        qk_norm : bool, optional
            Whether to use query and key normalization, default False
        cpu_offload : bool, optional
            Whether to offload processing to CPU, by default False
        layer_kernels : DotDict, optional
            A dict of layer implementations e.g. layer_kernels.Linear = "torch.nn.Linear"
            Defined in config/models/<model>.yaml
        shard_strategy : str, optional
            Strategy to shard tensors, by default "edges"
        graph_attention_backend: str, by default "triton"
            Backend to use for graph transformer conv, options are "triton" and "pyg"
        edge_pre_mlp: bool, by default False
            Allow for edge feature mixing
        """
        super().__init__(
            in_channels_src=in_channels_src,
            in_channels_dst=in_channels_dst,
            hidden_dim=hidden_dim,
            out_channels_dst=out_channels_dst,
            num_chunks=num_chunks,
            cpu_offload=cpu_offload,
            layer_kernels=layer_kernels,
        )

        self.num_chunks = num_chunks

        Linear = self.layer_factory.Linear

        self.proc = GraphTransformerMapperBlock(
            in_channels=hidden_dim,
            hidden_dim=mlp_hidden_ratio * hidden_dim,
            out_channels=hidden_dim,
            num_heads=num_heads,
            edge_dim=edge_dim,
            qk_norm=qk_norm,
            layer_kernels=self.layer_factory,
            shard_strategy=shard_strategy,
            graph_attention_backend=graph_attention_backend,
            edge_pre_mlp=edge_pre_mlp,
        )

        self.offload_layers(cpu_offload)

        self.emb_nodes_dst = Linear(self.in_channels_dst, self.hidden_dim)

        self.shard_strategy = shard_strategy

        assert shard_strategy in ["heads", "edges"], (
            f"Invalid shard strategy '{shard_strategy}' for {self.__class__.__name__}. "
            f"Supported strategies are 'heads' and 'edges'."
        )

    def prepare_edge_sharding_wrapper(
        self,
        x: PairTensor,
        shard_shapes: tuple[list[list[int]], list[list[int]]],
        batch_size: int,
        edge_attr: Tensor,
        edge_index: Adj,
        model_comm_group: Optional[ProcessGroup] = None,
        x_src_is_sharded: bool = False,
        x_dst_is_sharded: bool = False,
        cond: Optional[tuple[Tensor, Tensor]] = None,
        edge_shard_shapes: Optional[tuple] = None,
    ):
        x_src, x_dst = x
        shapes_src, shapes_dst = shard_shapes

        shapes_x_src = change_channels_in_shape(shapes_src, x_src.shape[-1])
        # gather/scatter if x_src is sharded, always reduce gradients in bwds
        x_src = sync_tensor(x_src, 0, shapes_x_src, model_comm_group, gather_in_fwd=x_src_is_sharded)

        if edge_shard_shapes is not None:
            # Edges are 1-hop sorted and sharded by graph provider
            shapes_edge_attr = edge_shard_shapes[0]
            shapes_edge_idx = edge_shard_shapes[1]
        else:
            # Edges not pre-sharded, do 1-hop sorting and sharding here
            src_size = sum(shape[0] for shape in shapes_src)
            dst_size = sum(shape[0] for shape in shapes_dst)
            edge_attr, edge_index, (shapes_edge_attr, shapes_edge_idx) = shard_edges_1hop(
                edge_attr, edge_index, src_size, dst_size, model_comm_group
            )

        # Relabel destination indices from global to local
        if model_comm_group is not None and model_comm_group.size() > 1:
            rank = model_comm_group.rank()
            dst_offset = sum(shapes_dst[i][0] for i in range(rank))
            edge_index = edge_index.clone()  # no in-place modification of pre-sharded tensor
            edge_index[1] -= dst_offset

        # at this point, x_src is synced i.e. full, x_dst is sharded, edges are sharded (incoming edges to x_dst)
        size_src_full_dst_shard = (x_src.shape[0], x_dst.shape[0])
        x_src, edge_index, nodes_src = drop_unconnected_src_nodes(x_src, edge_index, size_src_full_dst_shard)

        if cond is not None:  # sync cond_src to match x_src:
            cond_src, cond_dst = cond
            shapes_cond_src = change_channels_in_shape(shapes_src, cond_src.shape[-1])
            cond_src_full = sync_tensor(cond_src, 0, shapes_cond_src, model_comm_group, gather_in_fwd=True)
            cond = (cond_src_full[nodes_src], cond_dst)

        if not x_dst_is_sharded:
            x_dst = shard_tensor(x_dst, 0, shapes_dst, model_comm_group)

        return x_src, x_dst, edge_attr, edge_index, shapes_src, shapes_dst, cond

    def run_processor_chunk_edge_sharding(
        self,
        x: tuple[Tensor, Tensor],
        dst_chunk: Tensor,
        edge_attr: Tensor,
        edge_index: Adj,
        shapes: tuple[tuple[int], tuple[int]],
        batch_size: int,
        size: tuple[int],
        model_comm_group: Optional[ProcessGroup] = None,
        cond: Optional[tuple[Tensor, Tensor]] = None,
        **kwargs,
    ) -> Tensor:
        x_src, x_dst = x

        # get subgraph of x_dst_chunk and incoming edges, drop unconnected src nodes
        nodes_src_full = torch.arange(size[0], device=edge_index.device)
        edge_index, edge_attr = bipartite_subgraph(
            (nodes_src_full, dst_chunk),
            edge_index,
            edge_attr,
            size=size,
            relabel_nodes=True,
        )

        # drop unconnected src nodes and relabel edges
        x_src_chunk, edge_index_chunk, connected_src_nodes = drop_unconnected_src_nodes(x_src, edge_index, size)
        x_dst_chunk = x_dst[dst_chunk]
        chunk_size = (x_src_chunk.shape[0], x_dst_chunk.shape[0])

        if cond is not None:  # update cond with correct conditioning
            cond_src, cond_dst = cond
            cond = (cond_src[connected_src_nodes], cond_dst[dst_chunk])

        # pre-process chunk, embedding x_src/x_dst if not already done
        x_src_chunk, x_dst_chunk, _, _ = self.pre_process(
            x=(x_src_chunk, x_dst_chunk),
            shard_shapes=shapes,
            model_comm_group=model_comm_group,
            x_src_is_sharded=True,
            x_dst_is_sharded=True,
        )

        (_, x_dst_out), _ = self.proc(
            (x_src_chunk, x_dst_chunk),
            edge_attr,
            edge_index_chunk,
            shapes,
            batch_size,
            chunk_size,
            model_comm_group,
            cond=cond,
            **kwargs,
        )

        return self.post_process(
            x_dst=x_dst_out, shapes_dst=shapes[1], model_comm_group=model_comm_group, keep_x_dst_sharded=True
        )

    def mapper_forward_with_edge_sharding(
        self,
        x: PairTensor,
        batch_size: int,
        shard_shapes: tuple[list[list[int]], list[list[int]]],
        edge_attr: Tensor,
        edge_index: Adj,
        model_comm_group: Optional[ProcessGroup] = None,
        x_src_is_sharded: bool = False,
        x_dst_is_sharded: bool = False,
        keep_x_dst_sharded: bool = False,
        cond: Optional[tuple[Tensor, Tensor]] = None,
        edge_shard_shapes: Optional[tuple] = None,
        **kwargs,
    ) -> PairTensor:
        x_src, x_dst, edge_attr, edge_index, shapes_src, shapes_dst, cond = maybe_checkpoint(
            self.prepare_edge_sharding_wrapper,
            self.gradient_checkpointing,
            x,
            shard_shapes,
            batch_size,
            edge_attr,
            edge_index,
            model_comm_group,
            x_src_is_sharded,
            x_dst_is_sharded,
            cond,
            edge_shard_shapes,
        )

        size = (x_src.shape[0], x_dst.shape[0])  # node sizes of local graph shard
        num_chunks = max(self.num_chunks, NUM_CHUNKS_INFERENCE_MAPPER)

        dst_chunks = torch.arange(size[1], device=x_dst.device).tensor_split(num_chunks)
        out_channels = self.out_channels_dst if self.out_channels_dst is not None else self.hidden_dim
        out_type = torch.get_autocast_gpu_dtype() if torch.is_autocast_enabled() else x_dst.dtype
        out_dst = torch.empty((*x_dst.shape[:-1], out_channels), device=x_dst.device, dtype=out_type)

        for dst_chunk in dst_chunks:
            out_dst[dst_chunk] = maybe_checkpoint(
                self.run_processor_chunk_edge_sharding,
                self.gradient_checkpointing,
                (x_src, x_dst),
                dst_chunk,
                edge_attr,
                edge_index,
                (shapes_src, shapes_dst),
                batch_size,
                size,
                model_comm_group,
                cond,
                **kwargs,
            ).to(dtype=out_type)

        if not keep_x_dst_sharded:  # gather after processing chunks
            out_dst = gather_tensor(out_dst, 0, change_channels_in_shape(shapes_dst, out_channels), model_comm_group)

        return out_dst

    def mapper_forward_with_heads_sharding(
        self,
        x: PairTensor,
        batch_size: int,
        shard_shapes: tuple[list[list[int]], list[list[int]]],
        edge_attr: Tensor,
        edge_index: Adj,
        model_comm_group: Optional[ProcessGroup] = None,
        x_src_is_sharded: bool = False,
        x_dst_is_sharded: bool = False,
        keep_x_dst_sharded: bool = False,
        edge_shard_shapes: Optional[tuple] = None,
        **kwargs,
    ) -> PairTensor:
        size = (sum(x[0] for x in shard_shapes[0]), sum(x[0] for x in shard_shapes[1]))

        if edge_shard_shapes is not None:
            # Heads sharding needs full edge_index
            shapes_edge_attr = edge_shard_shapes[0]
            edge_index = gather_tensor(edge_index, 1, edge_shard_shapes[1], model_comm_group)
        else:
            shapes_edge_attr = get_shard_shapes(edge_attr, 0, model_comm_group)
            edge_attr = shard_tensor(edge_attr, 0, shapes_edge_attr, model_comm_group)

        x_src, x_dst, shapes_src, shapes_dst = self.pre_process(
            x=x,
            shard_shapes=shard_shapes,
            model_comm_group=model_comm_group,
            x_src_is_sharded=x_src_is_sharded,
            x_dst_is_sharded=x_dst_is_sharded,
        )

        (x_src, x_dst), edge_attr = self.proc(
            x=(x_src, x_dst),
            edge_attr=edge_attr,
            edge_index=edge_index,
            shapes=(shapes_src, shapes_dst, shapes_edge_attr),
            batch_size=batch_size,
            size=size,
            model_comm_group=model_comm_group,
            **kwargs,
        )

        x_dst = self.post_process(
            x_dst=x_dst, shapes_dst=shapes_dst, model_comm_group=model_comm_group, keep_x_dst_sharded=keep_x_dst_sharded
        )

        return x_dst

    def forward(
        self,
        x: PairTensor,
        batch_size: int,
        shard_shapes: tuple[list[list[int]], list[list[int]]],
        edge_attr: Tensor,
        edge_index: Adj,
        model_comm_group: Optional[ProcessGroup] = None,
        x_src_is_sharded: bool = False,
        x_dst_is_sharded: bool = False,
        keep_x_dst_sharded: bool = False,
        edge_shard_shapes: Optional[tuple] = None,
        **kwargs,
    ) -> PairTensor:

        kwargs_forward = {
            "x": x,
            "batch_size": batch_size,
            "shard_shapes": shard_shapes,
            "edge_attr": edge_attr,
            "edge_index": edge_index,
            "model_comm_group": model_comm_group,
            "x_src_is_sharded": x_src_is_sharded,
            "x_dst_is_sharded": x_dst_is_sharded,
            "keep_x_dst_sharded": keep_x_dst_sharded,
            "edge_shard_shapes": edge_shard_shapes,
            **kwargs,
        }

        if self.shard_strategy == "edges":
            return self.mapper_forward_with_edge_sharding(**kwargs_forward)
        else:  # self.shard_strategy == "heads"
            return maybe_checkpoint(
                self.mapper_forward_with_heads_sharding,
                self.gradient_checkpointing,
                **kwargs_forward,
            )


class GraphTransformerForwardMapper(GraphTransformerBaseMapper):
    """Graph Transformer Mapper from data -> hidden."""

    def __init__(
        self,
        *,
        in_channels_src: int,
        in_channels_dst: int,
        hidden_dim: int,
        num_chunks: int,
        num_heads: int,
        mlp_hidden_ratio: int,
        edge_dim: int,
        qk_norm: bool = False,
        cpu_offload: bool = False,
        layer_kernels: DotDict = None,
        shard_strategy: str = "edges",
        graph_attention_backend: str = "triton",
        edge_pre_mlp: bool = False,
        **kwargs,
    ) -> None:
        """Initialize GraphTransformerForwardMapper.

        Parameters
        ----------
        in_channels_src : int
            Input channels of the source node
        in_channels_dst : int
            Input channels of the destination node
        hidden_dim : int
            Hidden dimension
        num_chunks : int
            Number of chunks to split into
        num_heads: int
            Number of heads in transformer
        mlp_hidden_ratio: int
            ratio of mlp hidden dimension to embedding dimension
        edge_dim : int
            Edge feature dimension
        qk_norm : bool, optional
            Whether to use query and key normalization, default False
        cpu_offload : bool
            Whether to offload processing to CPU, by default False
        layer_kernels : DotDict
            A dict of layer implementations e.g. layer_kernels.Linear = "torch.nn.Linear"
        shard_strategy : str, optional
            Strategy to shard tensors, by default "edges"
        graph_attention_backend: str, by default "triton"
            Backend to use for graph transformer conv, options are "triton" and "pyg"
        edge_pre_mlp: bool, by default False
            Allow for edge feature mixing
        """
        super().__init__(
            in_channels_src=in_channels_src,
            in_channels_dst=in_channels_dst,
            hidden_dim=hidden_dim,
            out_channels_dst=None,
            num_chunks=num_chunks,
            cpu_offload=cpu_offload,
            qk_norm=qk_norm,
            num_heads=num_heads,
            mlp_hidden_ratio=mlp_hidden_ratio,
            edge_dim=edge_dim,
            layer_kernels=layer_kernels,
            shard_strategy=shard_strategy,
            graph_attention_backend=graph_attention_backend,
            edge_pre_mlp=edge_pre_mlp,
        )

        self.emb_nodes_src = self.layer_factory.Linear(self.in_channels_src, self.hidden_dim)

    def pre_process(self, x, shard_shapes, model_comm_group=None, x_src_is_sharded=False, x_dst_is_sharded=False):
        shapes_src, shapes_dst = shard_shapes
        x_src, x_dst = x
        if not x_src_is_sharded:
            x_src = shard_tensor(x_src, 0, shapes_src, model_comm_group)
        if not x_dst_is_sharded:
            x_dst = shard_tensor(x_dst, 0, shapes_dst, model_comm_group)
        x_src = self.emb_nodes_src(x_src)
        x_dst = self.emb_nodes_dst(x_dst)
        shapes_src = change_channels_in_shape(shapes_src, self.hidden_dim)
        shapes_dst = change_channels_in_shape(shapes_dst, self.hidden_dim)
        return x_src, x_dst, shapes_src, shapes_dst

    def post_process(self, x_dst, **kwargs):
        return x_dst

    def forward(
        self,
        x: PairTensor,
        batch_size: int,
        shard_shapes: tuple[list[list[int]], list[list[int]]],
        edge_attr: Tensor,
        edge_index: Adj,
        model_comm_group: Optional[ProcessGroup] = None,
        x_src_is_sharded: bool = False,
        x_dst_is_sharded: bool = False,
        keep_x_dst_sharded: bool = True,
        edge_shard_shapes: Optional[tuple] = None,
        **kwargs,
    ) -> PairTensor:
        x_dst = super().forward(
            x,
            batch_size,
            shard_shapes,
            edge_attr,
            edge_index,
            model_comm_group,
            x_src_is_sharded,
            x_dst_is_sharded,
            keep_x_dst_sharded,
            edge_shard_shapes,
            **kwargs,
        )
        return x[0], x_dst


class GraphTransformerBackwardMapper(GraphTransformerBaseMapper):
    """Graph Transformer Mapper from hidden -> data."""

    def __init__(
        self,
        *,
        in_channels_src: int,
        in_channels_dst: int,
        hidden_dim: int,
        out_channels_dst: Optional[int] = None,
        num_chunks: int,
        num_heads: int,
        mlp_hidden_ratio: int,
        edge_dim: int,
        qk_norm: bool = False,
        initialise_data_extractor_zero: bool = False,
        cpu_offload: bool = False,
        layer_kernels: DotDict = None,
        shard_strategy: str = "edges",
        graph_attention_backend: str = "triton",
        edge_pre_mlp: bool = False,
        **kwargs,
    ) -> None:
        """Initialize GraphTransformerBackwardMapper.

        Parameters
        ----------
        in_channels_src : int
            Input channels of the source node
        in_channels_dst : int
            Input channels of the destination node
        hidden_dim : int
            Hidden dimension
        out_channels_dst : int
            Output channels of the destination node
        num_chunks : int
            Number of chunks to split into
        num_heads: int
            Number of heads in transformer
        mlp_hidden_ratio: int
            Ratio of mlp hidden dimension to embedding dimension
        edge_dim : int
            Edge feature dimension
        qk_norm : bool, optional
            Whether to use query and key normalization, default False
        initialise_data_extractor_zero : bool, default False:
            Whether to initialise the data extractor to zero
        cpu_offload : bool, optional
            Whether to offload processing to CPU, by default False
        layer_kernels : DotDict, optional
            A dict of layer implementations e.g. layer_kernels.Linear = "torch.nn.Linear"
            Defined in config/models/<model>.yaml
        shard_strategy : str, optional
            Strategy to shard tensors, by default "edges"
        graph_attention_backend: str, by default "triton"
            Backend to use for graph transformer conv, options are "triton" and "pyg"
        edge_pre_mlp: bool, by default False
            Allow for edge feature mixing
        """
        super().__init__(
            in_channels_src=in_channels_src,
            in_channels_dst=in_channels_dst,
            hidden_dim=hidden_dim,
            out_channels_dst=out_channels_dst,
            num_chunks=num_chunks,
            cpu_offload=cpu_offload,
            qk_norm=qk_norm,
            num_heads=num_heads,
            mlp_hidden_ratio=mlp_hidden_ratio,
            edge_dim=edge_dim,
            layer_kernels=layer_kernels,
            shard_strategy=shard_strategy,
            graph_attention_backend=graph_attention_backend,
            edge_pre_mlp=edge_pre_mlp,
        )

        self.node_data_extractor = nn.Sequential(
            nn.LayerNorm(self.hidden_dim), nn.Linear(self.hidden_dim, self.out_channels_dst)
        )
        if initialise_data_extractor_zero:
            for module in self.node_data_extractor.modules():
                if isinstance(module, nn.Linear):
                    nn.init.constant_(module.weight, 0.0)
                    if module.bias is not None:
                        nn.init.constant_(module.bias, 0.0)

    def pre_process(self, x, shard_shapes, model_comm_group=None, x_src_is_sharded=False, x_dst_is_sharded=False):
        shapes_src, shapes_dst = shard_shapes
        x_src, x_dst = x
        shapes_src = change_channels_in_shape(shapes_src, self.hidden_dim)
        if not x_dst_is_sharded:
            x_dst = shard_tensor(x_dst, 0, shapes_dst, model_comm_group)
        x_dst = self.emb_nodes_dst(x_dst)
        shapes_dst = change_channels_in_shape(shapes_dst, self.hidden_dim)
        return x_src, x_dst, shapes_src, shapes_dst

    def post_process(self, x_dst, shapes_dst, model_comm_group=None, keep_x_dst_sharded=False):
        x_dst = self.node_data_extractor(x_dst)
        if not keep_x_dst_sharded:
            x_dst = gather_tensor(
                x_dst, 0, change_channels_in_shape(shapes_dst, self.out_channels_dst), model_comm_group
            )
        return x_dst


class GNNBaseMapper(BaseMapper, ABC):
    """Base for Graph Neural Network Mapper from hidden -> data or data -> hidden."""

    def __init__(
        self,
        *,
        in_channels_src: int,
        in_channels_dst: int,
        hidden_dim: int,
        out_channels_dst: Optional[int] = None,
        num_chunks: int,
        mlp_extra_layers: int,
        edge_dim: int,
        cpu_offload: bool = False,
        layer_kernels: DotDict = None,
        **kwargs,
    ) -> None:
        """Initialize GNNBaseMapper.

        Parameters
        ----------
        in_channels_src : int
            Input channels of the source node
        in_channels_dst : int
            Input channels of the destination node
        hidden_dim : int
            Hidden dimension
        out_channels_dst : int, optional
            Output channels of the destination node
        num_chunks : int
            Number of chunks to split into
        mlp_extra_layers : int
            Number of extra layers in MLP
        edge_dim : int
            Edge feature dimension
        cpu_offload : bool, optional
            Whether to offload processing to CPU, by default False
        layer_kernels : DotDict, optional
            A dict of layer implementations e.g. layer_kernels.Linear = "torch.nn.Linear"
            Defined in config/models/<model>.yaml
        """
        super().__init__(
            in_channels_src=in_channels_src,
            in_channels_dst=in_channels_dst,
            hidden_dim=hidden_dim,
            out_channels_dst=out_channels_dst,
            num_chunks=num_chunks,
            cpu_offload=cpu_offload,
            layer_kernels=layer_kernels,
        )

        self.emb_edges = MLP(
            in_features=edge_dim,
            hidden_dim=hidden_dim,
            out_features=hidden_dim,
            layer_kernels=self.layer_factory,
            n_extra_layers=mlp_extra_layers,
        )

    def mapper_forward(
        self,
        x: PairTensor,
        batch_size: int,
        shard_shapes: tuple[list[list[int]], list[list[int]]],
        edge_attr: Tensor,
        edge_index: Adj,
        model_comm_group: Optional[ProcessGroup] = None,
        x_src_is_sharded: bool = False,
        x_dst_is_sharded: bool = False,
        keep_x_dst_sharded: bool = False,
        edge_shard_shapes: Optional[tuple] = None,
        **kwargs,
    ) -> PairTensor:

        size = (sum(x[0] for x in shard_shapes[0]), sum(x[0] for x in shard_shapes[1]))

        if edge_shard_shapes is None:
            # Edges not pre-sharded, do 1-hop sorting and sharding here
            edge_attr, edge_index, _ = shard_edges_1hop(edge_attr, edge_index, size[0], size[1], model_comm_group)

        edge_attr = self.emb_edges(edge_attr)

        x_src, x_dst, shapes_src, shapes_dst = self.pre_process(
            x=x,
            shard_shapes=shard_shapes,
            model_comm_group=model_comm_group,
            x_src_is_sharded=x_src_is_sharded,
            x_dst_is_sharded=x_dst_is_sharded,
        )

        (x_src, x_dst), edge_attr = self.proc(
            (x_src, x_dst),
            edge_attr,
            edge_index,
            (shapes_src, shapes_dst),
            model_comm_group,
            size=size,
            **kwargs,
        )

        x_dst = self.post_process(
            x_dst=x_dst, shapes_dst=shapes_dst, model_comm_group=model_comm_group, keep_x_dst_sharded=keep_x_dst_sharded
        )

        return x_src, x_dst

    def forward(
        self,
        x: PairTensor,
        batch_size: int,
        shard_shapes: tuple[list[list[int]], list[list[int]]],
        edge_attr: Tensor,
        edge_index: Adj,
        model_comm_group: Optional[ProcessGroup] = None,
        x_src_is_sharded: bool = False,
        x_dst_is_sharded: bool = False,
        keep_x_dst_sharded: bool = False,
        edge_shard_shapes: Optional[tuple] = None,
        **kwargs,
    ) -> PairTensor:
        return maybe_checkpoint(
            self.mapper_forward,
            self.gradient_checkpointing,
            x=x,
            batch_size=batch_size,
            shard_shapes=shard_shapes,
            edge_attr=edge_attr,
            edge_index=edge_index,
            model_comm_group=model_comm_group,
            x_src_is_sharded=x_src_is_sharded,
            x_dst_is_sharded=x_dst_is_sharded,
            keep_x_dst_sharded=keep_x_dst_sharded,
            edge_shard_shapes=edge_shard_shapes,
            **kwargs,
        )


class GNNForwardMapper(GNNBaseMapper):
    """Graph Neural Network Mapper data -> hidden."""

    def __init__(
        self,
        *,
        in_channels_src: int,
        in_channels_dst: int,
        hidden_dim: int,
        out_channels_dst: Optional[int] = None,
        num_chunks: int,
        mlp_extra_layers: int,
        edge_dim: int,
        cpu_offload: bool = False,
        layer_kernels: DotDict,
        **kwargs,
    ) -> None:
        """Initialize GNNForwardMapper.

        Parameters
        ----------
        in_channels_src : int
            Input channels of the source node
        in_channels_dst : int
            Input channels of the destination node
        hidden_dim : int
            Hidden dimension
        out_channels_dst : int
            Output channels of the destination node, by default None
        num_chunks: int
            Number of chunks to split into
        mlp_extra_layers : int
            Number of extra layers in MLP
        edge_dim : int
            Edge feature dimension
        cpu_offload : bool, optional
            Whether to offload processing to CPU, by default False
        layer_kernels : DotDict
            A dict of layer implementations e.g. layer_kernels.Linear = "torch.nn.Linear"
            Defined in config/models/<model>.yaml
        """
        super().__init__(
            in_channels_src=in_channels_src,
            in_channels_dst=in_channels_dst,
            hidden_dim=hidden_dim,
            out_channels_dst=out_channels_dst,
            num_chunks=num_chunks,
            cpu_offload=cpu_offload,
            mlp_extra_layers=mlp_extra_layers,
            edge_dim=edge_dim,
            layer_kernels=layer_kernels,
        )

        self.proc = GraphConvMapperBlock(
            in_channels=hidden_dim,
            out_channels=hidden_dim,
            layer_kernels=self.layer_factory,
            mlp_extra_layers=mlp_extra_layers,
            update_src_nodes=True,
            num_chunks=num_chunks,
        )

        self.offload_layers(cpu_offload)

        self.emb_nodes_src = MLP(
            in_features=in_channels_src,
            hidden_dim=hidden_dim,
            out_features=hidden_dim,
            layer_kernels=self.layer_factory,
            n_extra_layers=mlp_extra_layers,
        )

        self.emb_nodes_dst = MLP(
            in_features=in_channels_dst,
            hidden_dim=hidden_dim,
            out_features=hidden_dim,
            layer_kernels=self.layer_factory,
            n_extra_layers=mlp_extra_layers,
        )

    def pre_process(self, x, shard_shapes, model_comm_group=None, x_src_is_sharded=False, x_dst_is_sharded=False):
        shapes_src, shapes_dst = shard_shapes
        x_src, x_dst = x
        if not x_src_is_sharded:
            x_src = shard_tensor(x_src, 0, shapes_src, model_comm_group)
        if not x_dst_is_sharded:
            x_dst = shard_tensor(x_dst, 0, shapes_dst, model_comm_group)
        x_src = self.emb_nodes_src(x_src)
        x_dst = self.emb_nodes_dst(x_dst)
        shapes_src = change_channels_in_shape(shapes_src, self.hidden_dim)
        shapes_dst = change_channels_in_shape(shapes_dst, self.hidden_dim)
        return x_src, x_dst, shapes_src, shapes_dst

    def post_process(self, x_dst, **kwargs):
        return x_dst


class GNNBackwardMapper(GNNBaseMapper):
    """Graph Neural Network Mapper from hidden -> data."""

    def __init__(
        self,
        *,
        in_channels_src: int,
        in_channels_dst: int,
        hidden_dim: int,
        out_channels_dst: Optional[int] = None,
        num_chunks: int,
        mlp_extra_layers: int,
        edge_dim: int,
        cpu_offload: bool = False,
        layer_kernels: DotDict,
        **kwargs,
    ) -> None:
        """Initialize GNNBackwardMapper.

        Parameters
        ----------
        in_channels_src : int
            Input channels of the source node
        in_channels_dst : int
            Input channels of the destination node
        hidden_dim : int
            Hidden dimension
        out_channels_dst : int
            Output channels of the destination node
        num_chunks: int
            Number of chunks to split into
        mlp_extra_layers : int
            Number of extra layers in MLP
        edge_dim : int
            Edge feature dimension
        cpu_offload : bool
            Whether to offload processing to CPU, by default False
        layer_kernels : DotDict
            A dict of layer implementations e.g. layer_kernels.Linear = "torch.nn.Linear"
            Defined in config/models/<model>.yaml
        """
        super().__init__(
            in_channels_src=in_channels_src,
            in_channels_dst=in_channels_dst,
            hidden_dim=hidden_dim,
            out_channels_dst=out_channels_dst,
            num_chunks=num_chunks,
            cpu_offload=cpu_offload,
            mlp_extra_layers=mlp_extra_layers,
            edge_dim=edge_dim,
            layer_kernels=layer_kernels,
        )

        self.proc = GraphConvMapperBlock(
            in_channels=hidden_dim,
            out_channels=hidden_dim,
            layer_kernels=self.layer_factory,
            mlp_extra_layers=mlp_extra_layers,
            update_src_nodes=False,
            num_chunks=num_chunks,
        )

        self.offload_layers(cpu_offload)

        self.node_data_extractor = MLP(
            in_features=self.hidden_dim,
            hidden_dim=self.hidden_dim,
            out_features=self.out_channels_dst,
            layer_kernels=self.layer_factory,
            n_extra_layers=mlp_extra_layers,
            layer_norm=False,
            final_activation=False,
        )

    def pre_process(self, x, shard_shapes, model_comm_group=None, x_src_is_sharded=False, x_dst_is_sharded=False):
        shapes_src, shapes_dst = shard_shapes
        x_src, x_dst = x
        shapes_src = change_channels_in_shape(shapes_src, self.hidden_dim)
        shapes_dst = change_channels_in_shape(shapes_dst, self.hidden_dim)
        return x_src, x_dst, shapes_src, shapes_dst

    def post_process(self, x_dst, shapes_dst, model_comm_group=None, keep_x_dst_sharded=False):
        x_dst = self.node_data_extractor(x_dst)
        if not keep_x_dst_sharded:
            x_dst = gather_tensor(
                x_dst, 0, change_channels_in_shape(shapes_dst, self.out_channels_dst), model_comm_group
            )
        return x_dst

    def forward(
        self,
        x: PairTensor,
        batch_size: int,
        shard_shapes: tuple[list[list[int]], list[list[int]]],
        edge_attr: Tensor,
        edge_index: Adj,
        model_comm_group: Optional[ProcessGroup] = None,
        x_src_is_sharded: bool = False,
        x_dst_is_sharded: bool = False,
        keep_x_dst_sharded: bool = False,
        edge_shard_shapes: Optional[tuple] = None,
        **kwargs,
    ) -> Tensor:

        _, x_dst = super().forward(
            x,
            batch_size,
            shard_shapes,
            edge_attr,
            edge_index,
            model_comm_group,
            x_src_is_sharded,
            x_dst_is_sharded,
            keep_x_dst_sharded,
            edge_shard_shapes,
            **kwargs,
        )
        return x_dst


class TransformerBaseMapper(BaseMapper, ABC):
    """Transformer Base Mapper from hidden -> data or data -> hidden."""

    def __init__(
        self,
        *,
        in_channels_src: int,
        in_channels_dst: int,
        hidden_dim: int,
        out_channels_dst: Optional[int] = None,
        num_chunks: int,
        num_heads: int,
        mlp_hidden_ratio: int,
        window_size: Optional[int] = None,
        dropout_p: float = 0.0,
        qk_norm: bool = False,
        attention_implementation: str = "flash_attention",
        softcap: Optional[float] = None,
        use_alibi_slopes: bool = False,
        use_rotary_embeddings: bool = False,
        cpu_offload: bool = False,
        layer_kernels: DotDict,
    ) -> None:
        """Initialize TransformerBaseMapper.

        Parameters
        ----------
        in_channels_src : int
            Input channels of the source node
        in_channels_dst : int
            Input channels of the destination node
        hidden_dim : int
            Hidden dimension
        out_channels_dst : int, optional
            Output channels of the destination node, by default None
        mlp_hidden_ratio: int
            Ratio of mlp hidden dimension to embedding dimension
        qk_norm: bool, optional
            Normalize query and key, by default False
        dropout_p: float, optional
            Dropout probability used for multi-head self attention, default 0.1
        attention_implementation: str
            A predefined string which selects which underlying attention
            implementation, by default "flash_attention"
        softcap : float, optional
            Anything > 0 activates softcapping flash attention, by default 0
        use_alibi_slopes : bool
            Use aLiBI option, only used for flash attention, by default False
        window_size: int, optional
            1/2 size of shifted window for attention computation, by default None
        cpu_offload : bool
            Whether to offload processing to CPU, by default False
        layer_kernels : DotDict
            A dict of layer implementations e.g. layer_kernels.Linear = "torch.nn.Linear"
            Defined in config/models/<model>.yaml
        """
        super().__init__(
            in_channels_src=in_channels_src,
            in_channels_dst=in_channels_dst,
            hidden_dim=hidden_dim,
            out_channels_dst=out_channels_dst,
            num_chunks=num_chunks,
            layer_kernels=layer_kernels,
            cpu_offload=cpu_offload,
        )

        self.proc = TransformerMapperBlock(
            num_channels=hidden_dim,
            hidden_dim=mlp_hidden_ratio * hidden_dim,
            num_heads=num_heads,
            window_size=window_size,
            layer_kernels=self.layer_factory,
            dropout_p=dropout_p,
            qk_norm=qk_norm,
            attention_implementation=attention_implementation,
            softcap=softcap,
            use_alibi_slopes=use_alibi_slopes,
            use_rotary_embeddings=use_rotary_embeddings,
        )

        self.offload_layers(cpu_offload)

        self.emb_nodes_dst = nn.Linear(self.in_channels_dst, self.hidden_dim)

    def mapper_forward(
        self,
        x: PairTensor,
        batch_size: int,
        shard_shapes: tuple[list[list[int]], list[list[int]]],
        model_comm_group: Optional[ProcessGroup] = None,
        x_src_is_sharded: bool = False,
        x_dst_is_sharded: bool = False,
        keep_x_dst_sharded: bool = False,
    ) -> PairTensor:

        x_src, x_dst, shapes_src, shapes_dst = self.pre_process(
            x=x,
            shard_shapes=shard_shapes,
            model_comm_group=model_comm_group,
            x_src_is_sharded=x_src_is_sharded,
            x_dst_is_sharded=x_dst_is_sharded,
        )

        (x_src, x_dst), _ = self.proc(
            (x_src, x_dst),
            (shapes_src, shapes_dst),
            batch_size,
            model_comm_group,
        )

        x_dst = self.post_process(
            x_dst=x_dst, shapes_dst=shapes_dst, model_comm_group=model_comm_group, keep_x_dst_sharded=keep_x_dst_sharded
        )

        return x_dst

    def forward(
        self,
        x: PairTensor,
        batch_size: int,
        shard_shapes: tuple[list[list[int]], list[list[int]]],
        edge_attr: Optional[Tensor] = None,
        edge_index: Optional[Adj] = None,
        model_comm_group: Optional[ProcessGroup] = None,
        x_src_is_sharded: bool = False,
        x_dst_is_sharded: bool = False,
        keep_x_dst_sharded: bool = False,
        edge_shard_shapes: Optional[tuple] = None,
        **kwargs,
    ) -> PairTensor:
        return maybe_checkpoint(
            self.mapper_forward,
            self.gradient_checkpointing,
            x=x,
            batch_size=batch_size,
            shard_shapes=shard_shapes,
            model_comm_group=model_comm_group,
            x_src_is_sharded=x_src_is_sharded,
            x_dst_is_sharded=x_dst_is_sharded,
            keep_x_dst_sharded=keep_x_dst_sharded,
            **kwargs,
        )


class TransformerForwardMapper(TransformerBaseMapper):
    """Transformer Mapper from data -> hidden."""

    def __init__(
        self,
        *,
        in_channels_src: int,
        in_channels_dst: int,
        hidden_dim: int,
        out_channels_dst: Optional[int] = None,
        num_chunks: int,
        num_heads: int,
        mlp_hidden_ratio: int,
        qk_norm: bool = False,
        dropout_p: float = 0.0,
        attention_implementation: str = "flash_attention",
        softcap: float = None,
        use_alibi_slopes: bool = False,
        cpu_offload: bool = False,
        window_size: Optional[int] = None,
        use_rotary_embeddings: bool = False,
        layer_kernels: DotDict,
        **kwargs,  # accept not needed extra arguments like subgraph etc.
    ) -> None:
        """Initialize TransformerForwardMapper.

        Parameters
        ----------
        in_channels_src : int
            Input channels of the source node
        in_channels_dst : int
            Input channels of the destination node
        hidden_dim : int
            Hidden dimension
        out_channels_dst : int, optional
            Output channels of the destination node, by default None
        mlp_hidden_ratio: int
            Ratio of mlp hidden dimension to embedding dimension
        qk_norm: bool, optional
            Normalize query and key, by default False
        dropout_p: float, optional
            Dropout probability used for multi-head self attention, default 0.1
        attention_implementation: str
            A predefined string which selects which underlying attention
            implementation, by default "flash_attention"
        softcap : float, optional
            Anything > 0 activates softcapping flash attention, by default 0
        use_alibi_slopes : bool
            Use aLiBI option, only used for flash attention, by default False
        window_size: int, optional
            1/2 size of shifted window for attention computation, by default None
        cpu_offload : bool
            Whether to offload processing to CPU, by default False
        layer_kernels : DotDict
            A dict of layer implementations e.g. layer_kernels.Linear = "torch.nn.Linear"
            Defined in config/models/<model>.yaml
        """
        super().__init__(
            in_channels_src=in_channels_src,
            in_channels_dst=in_channels_dst,
            hidden_dim=hidden_dim,
            layer_kernels=layer_kernels,
            out_channels_dst=out_channels_dst,
            num_chunks=num_chunks,
            cpu_offload=cpu_offload,
            num_heads=num_heads,
            mlp_hidden_ratio=mlp_hidden_ratio,
            window_size=window_size,
            dropout_p=dropout_p,
            qk_norm=qk_norm,
            attention_implementation=attention_implementation,
            softcap=softcap,
            use_alibi_slopes=use_alibi_slopes,
            use_rotary_embeddings=use_rotary_embeddings,
        )

        self.emb_nodes_src = nn.Linear(self.in_channels_src, self.hidden_dim)

    def pre_process(self, x, shard_shapes, model_comm_group=None, x_src_is_sharded=False, x_dst_is_sharded=False):
        shapes_src, shapes_dst = shard_shapes
        x_src, x_dst = x
        if not x_src_is_sharded:
            x_src = shard_tensor(x_src, 0, shapes_src, model_comm_group)
        if not x_dst_is_sharded:
            x_dst = shard_tensor(x_dst, 0, shapes_dst, model_comm_group)
        x_src = self.emb_nodes_src(x_src)
        x_dst = self.emb_nodes_dst(x_dst)
        shapes_src = change_channels_in_shape(shapes_src, self.hidden_dim)
        shapes_dst = change_channels_in_shape(shapes_dst, self.hidden_dim)
        return x_src, x_dst, shapes_src, shapes_dst

    def post_process(self, x_dst, **kwargs):
        return x_dst

    def forward(
        self,
        x: PairTensor,
        batch_size: int,
        shard_shapes: tuple[list[list[int]], list[list[int]], list[list[int]]],
        edge_attr: Optional[Tensor] = None,
        edge_index: Optional[Adj] = None,
        model_comm_group: Optional[ProcessGroup] = None,
        x_src_is_sharded: bool = False,
        x_dst_is_sharded: bool = False,
        keep_x_dst_sharded: bool = False,
        edge_shard_shapes: Optional[tuple] = None,
        **kwargs,
    ) -> PairTensor:
        x_dst = super().forward(
            x,
            batch_size,
            shard_shapes,
            edge_attr,
            edge_index,
            model_comm_group,
            x_src_is_sharded,
            x_dst_is_sharded,
            keep_x_dst_sharded,
            edge_shard_shapes,
            **kwargs,
        )
        return x[0], x_dst


class TransformerBackwardMapper(TransformerBaseMapper):
    """Graph Transformer Mapper from hidden -> data."""

    def __init__(
        self,
        *,
        in_channels_src: int,
        in_channels_dst: int,
        hidden_dim: int,
        out_channels_dst: Optional[int] = None,
        num_chunks: int,
        num_heads: int,
        mlp_hidden_ratio: int,
        qk_norm: bool = False,
        dropout_p: float = 0.0,
        attention_implementation: str = "flash_attention",
        softcap: float = None,
        use_alibi_slopes: bool = False,
        cpu_offload: bool = False,
        window_size: Optional[int] = None,
        use_rotary_embeddings: bool = False,
        layer_kernels: DotDict,
        **kwargs,  # accept not needed extra arguments like subgraph etc.
    ) -> None:
        """Initialize TransformerBackwardMapper.

        Parameters
        ----------
        in_channels_src : int
            Input channels of the source node
        in_channels_dst : int
            Input channels of the destination node
        hidden_dim : int
            Hidden dimension
        out_channels_dst : int, optional
            Output channels of the destination node, by default None
        mlp_hidden_ratio: int
            Ratio of mlp hidden dimension to embedding dimension
        qk_norm: bool, optional
            Normalize query and key, by default False
        dropout_p: float, optional
            Dropout probability used for multi-head self attention, default 0.1
        attention_implementation: str
            A predefined string which selects which underlying attention
            implementation, by default "flash_attention"
        softcap : float, optional
            Anything > 0 activates softcapping flash attention, by default 0
        use_alibi_slopes : bool
            Use aLiBI option, only used for flash attention, by default False
        window_size: int, optional
            1/2 size of shifted window for attention computation, by default None
        cpu_offload : bool
            Whether to offload processing to CPU, by default False
        layer_kernels : DotDict
            A dict of layer implementations e.g. layer_kernels.Linear = "torch.nn.Linear"
            Defined in config/models/<model>.yaml
        """
        super().__init__(
            in_channels_src=in_channels_src,
            in_channels_dst=in_channels_dst,
            hidden_dim=hidden_dim,
            layer_kernels=layer_kernels,
            out_channels_dst=out_channels_dst,
            num_chunks=num_chunks,
            cpu_offload=cpu_offload,
            num_heads=num_heads,
            mlp_hidden_ratio=mlp_hidden_ratio,
            window_size=window_size,
            dropout_p=dropout_p,
            qk_norm=qk_norm,
            attention_implementation=attention_implementation,
            softcap=softcap,
            use_alibi_slopes=use_alibi_slopes,
            use_rotary_embeddings=use_rotary_embeddings,
        )

        self.node_data_extractor = nn.Sequential(
            nn.LayerNorm(self.hidden_dim), nn.Linear(self.hidden_dim, self.out_channels_dst)
        )

    def pre_process(self, x, shard_shapes, model_comm_group=None, x_src_is_sharded=False, x_dst_is_sharded=False):
        shapes_src, shapes_dst = shard_shapes
        x_src, x_dst = x
        shapes_src = change_channels_in_shape(shapes_src, self.hidden_dim)
        if not x_dst_is_sharded:
            x_dst = shard_tensor(x_dst, 0, shapes_dst, model_comm_group)
        x_dst = self.emb_nodes_dst(x_dst)
        shapes_dst = change_channels_in_shape(shapes_dst, self.hidden_dim)
        return x_src, x_dst, shapes_src, shapes_dst

    def post_process(self, x_dst, shapes_dst, model_comm_group=None, keep_x_dst_sharded=False):
        x_dst = self.node_data_extractor(x_dst)
        if not keep_x_dst_sharded:
            x_dst = gather_tensor(
                x_dst, 0, change_channels_in_shape(shapes_dst, self.out_channels_dst), model_comm_group
            )
        return x_dst
