model/graph_cast_net.py

import logging
import warnings
from dataclasses import dataclass
from typing import Any, Optional

import torch
from torch import Tensor
from torch import nn
try:
    from typing import Self
except ImportError:
    # for Python versions < 3.11
    from typing_extensions import Self

from onescience.modules.decoder.mesh_graph_decoder import MeshGraphDecoder
from onescience.modules.embedding.graphcast_embedder import (
    GraphCastDecoderEmbedder,
    GraphCastEncoderEmbedder,
)
from onescience.modules.encoder.mesh_graph_encoder import MeshGraphEncoder
from onescience.modules.mlp.mesh_graph_mlp import MeshGraphMLP
from onescience.modules.utils.gnnlayer_utils import CuGraphCSC, set_checkpoint_fn
from onescience.modules.layer.activations import get_activation
from onescience.modules.utils.graphcast.graph import Graph

from onescience.models.meta import ModelMetaData

from .graph_cast_processor import (
    GraphCastProcessor,
    GraphCastProcessorGraphTransformer,
)

logger = logging.getLogger(__name__)


def get_lat_lon_partition_separators(partition_size: int):
    """Utility Function to get separation intervals for lat-lon
    grid for partition_sizes of interest.

    Parameters
    ----------
    partition_size : int
        size of graph partition
    """

    def _divide(num_lat_chunks: int, num_lon_chunks: int):
        # divide lat-lon grid into equally-sizes chunks along both latitude and longitude
        if (num_lon_chunks * num_lat_chunks) != partition_size:
            raise ValueError(
                "Can't divide lat-lon grid into grid {num_lat_chunks} x {num_lon_chunks} chunks for partition_size={partition_size}."
            )
        # divide latitutude into num_lat_chunks of size 180 / num_lat_chunks
        # divide longitude into chunks of size 360 / (partition_size / num_lat_chunks)
        lat_bin_width = 180.0 / num_lat_chunks
        lon_bin_width = 360.0 / num_lon_chunks

        lat_ranges = []
        lon_ranges = []

        for p_lat in range(num_lat_chunks):
            for p_lon in range(num_lon_chunks):
                lat_ranges += [
                    (lat_bin_width * p_lat - 90.0, lat_bin_width * (p_lat + 1) - 90.0)
                ]
                lon_ranges += [
                    (lon_bin_width * p_lon - 180.0, lon_bin_width * (p_lon + 1) - 180.0)
                ]

        lat_ranges[-1] = (lat_ranges[-1][0], None)
        lon_ranges[-1] = (lon_ranges[-1][0], None)

        return lat_ranges, lon_ranges

    # use two closest factors of partition_size
    lat_chunks, lon_chunks, i = 1, partition_size, 0
    while lat_chunks < lon_chunks:
        i += 1
        if partition_size % i == 0:
            lat_chunks = i
            lon_chunks = partition_size // lat_chunks

    lat_ranges, lon_ranges = _divide(lat_chunks, lon_chunks)

    # mainly for debugging
    if (lat_ranges is None) or (lon_ranges is None):
        raise ValueError("unexpected error, abort")

    min_seps = []
    max_seps = []

    for i in range(partition_size):
        lat = lat_ranges[i]
        lon = lon_ranges[i]
        min_seps.append([lat[0], lon[0]])
        max_seps.append([lat[1], lon[1]])

    return min_seps, max_seps


@dataclass
class MetaData(ModelMetaData):
    name: str = "GraphCastNet"
    # Optimization
    jit: bool = False
    cuda_graphs: bool = False
    amp_cpu: bool = False
    amp_gpu: bool = True
    torch_fx: bool = False
    # Data type
    bf16: bool = True
    # Inference
    onnx: bool = False
    # Physics informed
    func_torch: bool = False
    auto_grad: bool = False


class GraphCastNet(nn.Module):
    def __init__(
        self,
        mesh_level: Optional[int] = 6,
        multimesh_level: Optional[int] = None,
        multimesh: bool = True,
        input_res: tuple = (721, 1440),
        input_dim_grid_nodes: int = 237,
        input_dim_mesh_nodes: int = 3,
        input_dim_edges: int = 4,
        output_dim_grid_nodes: int = 227,
        processor_type: str = "MessagePassing",
        khop_neighbors: int = 32,
        num_attention_heads: int = 4,
        processor_layers: int = 16,
        hidden_layers: int = 1,
        hidden_dim: int = 512,
        aggregation: str = "sum",
        activation_fn: str = "silu",
        norm_type: str = "LayerNorm",
        use_cugraphops_encoder: bool = False,
        use_cugraphops_processor: bool = False,
        use_cugraphops_decoder: bool = False,
        do_concat_trick: bool = False,
        recompute_activation: bool = False,
        partition_size: int = 1,
        partition_group_name: Optional[str] = None,
        use_lat_lon_partitioning: bool = False,
        expect_partitioned_input: bool = False,
        global_features_on_rank_0: bool = False,
        produce_aggregated_output: bool = True,
        produce_aggregated_output_on_all_ranks: bool = True,
    ):
        super().__init__()

        # 'multimesh_level' deprecation handling
        if multimesh_level is not None:
            warnings.warn(
                "'multimesh_level' is deprecated and will be removed in a future version. Use 'mesh_level' instead.",
                DeprecationWarning,
                stacklevel=2,
            )
            mesh_level = multimesh_level

        self.processor_type = processor_type
        if self.processor_type == "MessagePassing":
            khop_neighbors = 0
        self.is_distributed = False
        if partition_size > 1:
            self.is_distributed = True
        self.expect_partitioned_input = expect_partitioned_input
        self.global_features_on_rank_0 = global_features_on_rank_0
        self.produce_aggregated_output = produce_aggregated_output
        self.produce_aggregated_output_on_all_ranks = (
            produce_aggregated_output_on_all_ranks
        )
        self.partition_group_name = partition_group_name

        # create the lat_lon_grid
        self.latitudes = torch.linspace(-90, 90, steps=input_res[0])
        self.longitudes = torch.linspace(-180, 180, steps=input_res[1] + 1)[1:]
        self.lat_lon_grid = torch.stack(
            torch.meshgrid(self.latitudes, self.longitudes, indexing="ij"), dim=-1
        )

        # Set activation function
        activation_fn = get_activation(activation_fn)

        # construct the graph
        self.graph = Graph(self.lat_lon_grid, mesh_level, multimesh, khop_neighbors)

        self.mesh_graph, self.attn_mask = self.graph.create_mesh_graph(verbose=False)
        self.g2m_graph = self.graph.create_g2m_graph(verbose=False)
        self.m2g_graph = self.graph.create_m2g_graph(verbose=False)

        self.g2m_edata = self.g2m_graph.edata["x"]
        self.m2g_edata = self.m2g_graph.edata["x"]
        self.mesh_ndata = self.mesh_graph.ndata["x"]
        if self.processor_type == "MessagePassing":
            self.mesh_edata = self.mesh_graph.edata["x"]
        elif self.processor_type == "GraphTransformer":
            # Dummy tensor to avoid breaking the API
            self.mesh_edata = torch.zeros((1, input_dim_edges))
        else:
            raise ValueError(f"Invalid processor type {processor_type}")

        if use_cugraphops_encoder or self.is_distributed:
            kwargs = {}
            if use_lat_lon_partitioning:
                min_seps, max_seps = get_lat_lon_partition_separators(partition_size)
                kwargs = {
                    "src_coordinates": self.g2m_graph.srcdata["lat_lon"],
                    "dst_coordinates": self.g2m_graph.dstdata["lat_lon"],
                    "coordinate_separators_min": min_seps,
                    "coordinate_separators_max": max_seps,
                }
            self.g2m_graph, edge_perm = CuGraphCSC.from_dgl(
                graph=self.g2m_graph,
                partition_size=partition_size,
                partition_group_name=partition_group_name,
                partition_by_bbox=use_lat_lon_partitioning,
                **kwargs,
            )
            self.g2m_edata = self.g2m_edata[edge_perm]

            if self.is_distributed:
                self.g2m_edata = self.g2m_graph.get_edge_features_in_partition(
                    self.g2m_edata
                )

        if use_cugraphops_decoder or self.is_distributed:
            kwargs = {}
            if use_lat_lon_partitioning:
                min_seps, max_seps = get_lat_lon_partition_separators(partition_size)
                kwargs = {
                    "src_coordinates": self.m2g_graph.srcdata["lat_lon"],
                    "dst_coordinates": self.m2g_graph.dstdata["lat_lon"],
                    "coordinate_separators_min": min_seps,
                    "coordinate_separators_max": max_seps,
                }

            self.m2g_graph, edge_perm = CuGraphCSC.from_dgl(
                graph=self.m2g_graph,
                partition_size=partition_size,
                partition_group_name=partition_group_name,
                partition_by_bbox=use_lat_lon_partitioning,
                **kwargs,
            )
            self.m2g_edata = self.m2g_edata[edge_perm]

            if self.is_distributed:
                self.m2g_edata = self.m2g_graph.get_edge_features_in_partition(
                    self.m2g_edata
                )

        if use_cugraphops_processor or self.is_distributed:
            kwargs = {}
            if use_lat_lon_partitioning:
                min_seps, max_seps = get_lat_lon_partition_separators(partition_size)
                kwargs = {
                    "src_coordinates": self.mesh_graph.ndata["lat_lon"],
                    "dst_coordinates": self.mesh_graph.ndata["lat_lon"],
                    "coordinate_separators_min": min_seps,
                    "coordinate_separators_max": max_seps,
                }

            self.mesh_graph, edge_perm = CuGraphCSC.from_dgl(
                graph=self.mesh_graph,
                partition_size=partition_size,
                partition_group_name=partition_group_name,
                partition_by_bbox=use_lat_lon_partitioning,
                **kwargs,
            )
            self.mesh_edata = self.mesh_edata[edge_perm]
            if self.is_distributed:
                self.mesh_edata = self.mesh_graph.get_edge_features_in_partition(
                    self.mesh_edata
                )
                self.mesh_ndata = self.mesh_graph.get_dst_node_features_in_partition(
                    self.mesh_ndata
                )

        self.input_dim_grid_nodes = input_dim_grid_nodes
        self.output_dim_grid_nodes = output_dim_grid_nodes
        self.input_res = input_res

        # by default: don't checkpoint at all
        self.model_checkpoint_fn = set_checkpoint_fn(False)
        self.encoder_checkpoint_fn = set_checkpoint_fn(False)
        self.decoder_checkpoint_fn = set_checkpoint_fn(False)

        # initial feature embedder
        self.encoder_embedder = GraphCastEncoderEmbedder(
            input_dim_grid_nodes=input_dim_grid_nodes,
            input_dim_mesh_nodes=input_dim_mesh_nodes,
            input_dim_edges=input_dim_edges,
            output_dim=hidden_dim,
            hidden_dim=hidden_dim,
            hidden_layers=hidden_layers,
            activation_fn=activation_fn,
            norm_type=norm_type,
            recompute_activation=recompute_activation,
        )
        self.decoder_embedder = GraphCastDecoderEmbedder(
            input_dim_edges=input_dim_edges,
            output_dim=hidden_dim,
            hidden_dim=hidden_dim,
            hidden_layers=hidden_layers,
            activation_fn=activation_fn,
            norm_type=norm_type,
            recompute_activation=recompute_activation,
        )

        # grid2mesh encoder
        self.encoder = MeshGraphEncoder(
            aggregation=aggregation,
            input_dim_src_nodes=hidden_dim,
            input_dim_dst_nodes=hidden_dim,
            input_dim_edges=hidden_dim,
            output_dim_src_nodes=hidden_dim,
            output_dim_dst_nodes=hidden_dim,
            output_dim_edges=hidden_dim,
            hidden_dim=hidden_dim,
            hidden_layers=hidden_layers,
            activation_fn=activation_fn,
            norm_type=norm_type,
            do_concat_trick=do_concat_trick,
            recompute_activation=recompute_activation,
        )

        # icosahedron processor
        if processor_layers <= 2:
            raise ValueError("Expected at least 3 processor layers")
        if processor_type == "MessagePassing":
            self.processor_encoder = GraphCastProcessor(
                aggregation=aggregation,
                processor_layers=1,
                input_dim_nodes=hidden_dim,
                input_dim_edges=hidden_dim,
                hidden_dim=hidden_dim,
                hidden_layers=hidden_layers,
                activation_fn=activation_fn,
                norm_type=norm_type,
                do_concat_trick=do_concat_trick,
                recompute_activation=recompute_activation,
            )
            self.processor = GraphCastProcessor(
                aggregation=aggregation,
                processor_layers=processor_layers - 2,
                input_dim_nodes=hidden_dim,
                input_dim_edges=hidden_dim,
                hidden_dim=hidden_dim,
                hidden_layers=hidden_layers,
                activation_fn=activation_fn,
                norm_type=norm_type,
                do_concat_trick=do_concat_trick,
                recompute_activation=recompute_activation,
            )
            self.processor_decoder = GraphCastProcessor(
                aggregation=aggregation,
                processor_layers=1,
                input_dim_nodes=hidden_dim,
                input_dim_edges=hidden_dim,
                hidden_dim=hidden_dim,
                hidden_layers=hidden_layers,
                activation_fn=activation_fn,
                norm_type=norm_type,
                do_concat_trick=do_concat_trick,
                recompute_activation=recompute_activation,
            )
        else:
            self.processor_encoder = torch.nn.Identity()
            self.processor = GraphCastProcessorGraphTransformer(
                attention_mask=self.attn_mask,
                num_attention_heads=num_attention_heads,
                processor_layers=processor_layers,
                input_dim_nodes=hidden_dim,
                hidden_dim=hidden_dim,
            )
            self.processor_decoder = torch.nn.Identity()

        # mesh2grid decoder
        self.decoder = MeshGraphDecoder(
            aggregation=aggregation,
            input_dim_src_nodes=hidden_dim,
            input_dim_dst_nodes=hidden_dim,
            input_dim_edges=hidden_dim,
            output_dim_dst_nodes=hidden_dim,
            output_dim_edges=hidden_dim,
            hidden_dim=hidden_dim,
            hidden_layers=hidden_layers,
            activation_fn=activation_fn,
            norm_type=norm_type,
            do_concat_trick=do_concat_trick,
            recompute_activation=recompute_activation,
        )

        # final MLP
        self.finale = MeshGraphMLP(
            input_dim=hidden_dim,
            output_dim=output_dim_grid_nodes,
            hidden_dim=hidden_dim,
            hidden_layers=hidden_layers,
            activation_fn=activation_fn,
            norm_type=None,
            recompute_activation=recompute_activation,
        )

    def set_checkpoint_model(self, checkpoint_flag: bool):
        """Sets checkpoint function for the entire model.

        This function returns the appropriate checkpoint function based on the
        provided `checkpoint_flag` flag. If `checkpoint_flag` is True, the
        function returns the checkpoint function from PyTorch's
        `torch.utils.checkpoint`. In this case, all the other gradient checkpoitings
        will be disabled. Otherwise, it returns an identity function
        that simply passes the inputs through the given layer.

        Parameters
        ----------
        checkpoint_flag : bool
            Whether to use checkpointing for gradient computation. Checkpointing
            can reduce memory usage during backpropagation at the cost of
            increased computation time.

        Returns
        -------
        Callable
            The selected checkpoint function to use for gradient computation.
        """
        # force a single checkpoint for the whole model
        self.model_checkpoint_fn = set_checkpoint_fn(checkpoint_flag)
        if checkpoint_flag:
            self.processor.set_checkpoint_segments(-1)
            self.encoder_checkpoint_fn = set_checkpoint_fn(False)
            self.decoder_checkpoint_fn = set_checkpoint_fn(False)

    def set_checkpoint_processor(self, checkpoint_segments: int):
        """Sets checkpoint function for the processor excluding the first and last
        layers.

        This function returns the appropriate checkpoint function based on the
        provided `checkpoint_segments` flag. If `checkpoint_segments` is positive,
        the function returns the checkpoint function from PyTorch's
        `torch.utils.checkpoint`, with number of checkpointing segments equal to
        `checkpoint_segments`. Otherwise, it returns an identity function
        that simply passes the inputs through the given layer.

        Parameters
        ----------
        checkpoint_segments : int
            Number of checkpointing segments for gradient computation. Checkpointing
            can reduce memory usage during backpropagation at the cost of
            increased computation time.

        Returns
        -------
        Callable
            The selected checkpoint function to use for gradient computation.
        """
        self.processor.set_checkpoint_segments(checkpoint_segments)

    def set_checkpoint_encoder(self, checkpoint_flag: bool):
        """Sets checkpoint function for the embedder, encoder, and the first of
        the processor.

        This function returns the appropriate checkpoint function based on the
        provided `checkpoint_flag` flag. If `checkpoint_flag` is True, the
        function returns the checkpoint function from PyTorch's
        `torch.utils.checkpoint`. Otherwise, it returns an identity function
        that simply passes the inputs through the given layer.

        Parameters
        ----------
        checkpoint_flag : bool
            Whether to use checkpointing for gradient computation. Checkpointing
            can reduce memory usage during backpropagation at the cost of
            increased computation time.

        Returns
        -------
        Callable
            The selected checkpoint function to use for gradient computation.
        """
        self.encoder_checkpoint_fn = set_checkpoint_fn(checkpoint_flag)

    def set_checkpoint_decoder(self, checkpoint_flag: bool):
        """Sets checkpoint function for the last layer of the processor, the decoder,
        and the final MLP.

        This function returns the appropriate checkpoint function based on the
        provided `checkpoint_flag` flag. If `checkpoint_flag` is True, the
        function returns the checkpoint function from PyTorch's
        `torch.utils.checkpoint`. Otherwise, it returns an identity function
        that simply passes the inputs through the given layer.

        Parameters
        ----------
        checkpoint_flag : bool
            Whether to use checkpointing for gradient computation. Checkpointing
            can reduce memory usage during backpropagation at the cost of
            increased computation time.

        Returns
        -------
        Callable
            The selected checkpoint function to use for gradient computation.
        """
        self.decoder_checkpoint_fn = set_checkpoint_fn(checkpoint_flag)

    def encoder_forward(
        self,
        grid_nfeat: Tensor,
    ) -> Tensor:
        """Forward method for the embedder, encoder, and the first of the processor.

        Parameters
        ----------
        grid_nfeat : Tensor
            Node features for the latitude-longitude grid.

        Returns
        -------
        mesh_efeat_processed: Tensor
            Processed edge features for the multimesh.
        mesh_nfeat_processed: Tensor
            Processed node features for the multimesh.
        grid_nfeat_encoded: Tensor
            Encoded node features for the latitude-longitude grid.
        """

        # embedd graph features
        (
            grid_nfeat_embedded,
            mesh_nfeat_embedded,
            g2m_efeat_embedded,
            mesh_efeat_embedded,
        ) = self.encoder_embedder(
            grid_nfeat,
            self.mesh_ndata,
            self.g2m_edata,
            self.mesh_edata
        )

        # encode lat/lon to multimesh
        grid_nfeat_encoded, mesh_nfeat_encoded = self.encoder(
            g2m_efeat_embedded,
            grid_nfeat_embedded,
            mesh_nfeat_embedded,
            self.g2m_graph,
        )

        # process multimesh graph
        if self.processor_type == "MessagePassing":
            mesh_efeat_processed, mesh_nfeat_processed = self.processor_encoder(
                mesh_efeat_embedded,
                mesh_nfeat_encoded,
                self.mesh_graph,
            )
        else:
            mesh_nfeat_processed = self.processor_encoder(
                mesh_nfeat_encoded,
            )
            mesh_efeat_processed = None
        return mesh_efeat_processed, mesh_nfeat_processed, grid_nfeat_encoded

    def decoder_forward(
        self,
        mesh_efeat_processed: Tensor,
        mesh_nfeat_processed: Tensor,
        grid_nfeat_encoded: Tensor,
    ) -> Tensor:
        """Forward method for the last layer of the processor, the decoder,
        and the final MLP.

        Parameters
        ----------
        mesh_efeat_processed : Tensor
            Multimesh edge features processed by the processor.
        mesh_nfeat_processed : Tensor
            Multi-mesh node features processed by the processor.
        grid_nfeat_encoded : Tensor
            The encoded node features for the latitude-longitude grid.

        Returns
        -------
        grid_nfeat_finale: Tensor
            The final node features for the latitude-longitude grid.
        """

        # process multimesh graph
        if self.processor_type == "MessagePassing":
            _, mesh_nfeat_processed = self.processor_decoder(
                mesh_efeat_processed,
                mesh_nfeat_processed,
                self.mesh_graph,
            )
        else:
            mesh_nfeat_processed = self.processor_decoder(
                mesh_nfeat_processed,
            )

        m2g_efeat_embedded = self.decoder_embedder(self.m2g_edata)

        # decode multimesh to lat/lon
        grid_nfeat_decoded = self.decoder(
            m2g_efeat_embedded, grid_nfeat_encoded, mesh_nfeat_processed, self.m2g_graph
        )

        # map to the target output dimension
        grid_nfeat_finale = self.finale(
            grid_nfeat_decoded,
        )

        return grid_nfeat_finale

    def custom_forward(self, grid_nfeat: Tensor) -> Tensor:
        """GraphCast forward method with support for gradient checkpointing.

        Parameters
        ----------
        grid_nfeat : Tensor
            Node features of the latitude-longitude graph.

        Returns
        -------
        grid_nfeat_finale: Tensor
            Predicted node features of the latitude-longitude graph.
        """
        (
            mesh_efeat_processed,
            mesh_nfeat_processed,
            grid_nfeat_encoded,
        ) = self.encoder_checkpoint_fn(
            self.encoder_forward,
            grid_nfeat,
            use_reentrant=False,
            preserve_rng_state=False,
        )

        # checkpoint of processor done in processor itself
        if self.processor_type == "MessagePassing":
            mesh_efeat_processed, mesh_nfeat_processed = self.processor(
                mesh_efeat_processed,
                mesh_nfeat_processed,
                self.mesh_graph,
            )
        else:
            mesh_nfeat_processed = self.processor(
                mesh_nfeat_processed,
            )
            mesh_efeat_processed = None

        grid_nfeat_finale = self.decoder_checkpoint_fn(
            self.decoder_forward,
            mesh_efeat_processed,
            mesh_nfeat_processed,
            grid_nfeat_encoded,
            use_reentrant=False,
            preserve_rng_state=False,
        )

        return grid_nfeat_finale

    def forward(
        self,
        grid_nfeat: Tensor,
    ) -> Tensor:
        invar = self.prepare_input(
            grid_nfeat, self.expect_partitioned_input, self.global_features_on_rank_0
        )
        outvar = self.model_checkpoint_fn(
            self.custom_forward,
            invar,
            use_reentrant=False,
            preserve_rng_state=False,
        )
        outvar = self.prepare_output(
            outvar,
            self.produce_aggregated_output,
            self.produce_aggregated_output_on_all_ranks,
        )
        return outvar

    def prepare_input(
        self,
        invar: Tensor,
        expect_partitioned_input: bool,
        global_features_on_rank_0: bool,
    ) -> Tensor:
        """Prepares the input to the model in the required shape.

        Parameters
        ----------
        invar : Tensor
            Input in the shape [N, C, H, W].

        expect_partitioned_input : bool
            flag indicating whether input is partioned according to graph partitioning scheme

        global_features_on_rank_0 : bool
            Flag indicating whether input is in its "global" form only on group_rank 0 which
            requires a scatter operation beforehand. Note that only either this flag or
            expect_partitioned_input can be set at a time.

        Returns
        -------
        Tensor
            Reshaped input.
        """
        if global_features_on_rank_0 and expect_partitioned_input:
            raise ValueError(
                "global_features_on_rank_0 and expect_partitioned_input cannot be set at the same time."
            )

        if not self.is_distributed:
            if invar.size(0) != 1:
                raise ValueError("GraphCast does not support batch size > 1")
            invar = invar[0].view(self.input_dim_grid_nodes, -1).permute(1, 0)

        else:
            # is_distributed
            if not expect_partitioned_input:
                # global_features_on_rank_0
                if invar.size(0) != 1:
                    raise ValueError("GraphCast does not support batch size > 1")

                invar = invar[0].view(self.input_dim_grid_nodes, -1).permute(1, 0)

                # scatter global features
                invar = self.g2m_graph.get_src_node_features_in_partition(
                    invar,
                    scatter_features=global_features_on_rank_0,
                )

        return invar

    def prepare_output(
        self,
        outvar: Tensor,
        produce_aggregated_output: bool,
        produce_aggregated_output_on_all_ranks: bool = True,
    ) -> Tensor:
        """Prepares the output of the model in the shape [N, C, H, W].

        Parameters
        ----------
        outvar : Tensor
            Output of the final MLP of the model.

        produce_aggregated_output : bool
            flag indicating whether output is gathered onto each rank
            or kept distributed

        produce_aggregated_output_on_all_ranks : bool
            flag indicating whether output is gatherered on each rank
            or only gathered at group_rank 0, True by default and
            only valid if produce_aggregated_output is set.

        Returns
        -------
        Tensor
            The reshaped output of the model.
        """
        if produce_aggregated_output or not self.is_distributed:
            # default case: output of shape [N, C, H, W]
            if self.is_distributed:
                outvar = self.m2g_graph.get_global_dst_node_features(
                    outvar,
                    get_on_all_ranks=produce_aggregated_output_on_all_ranks,
                )

            outvar = outvar.permute(1, 0)
            outvar = outvar.view(self.output_dim_grid_nodes, *self.input_res)
            outvar = torch.unsqueeze(outvar, dim=0)

        return outvar

    def to(self, *args: Any, **kwargs: Any) -> Self:
        """Moves the object to the specified device, dtype, or format.
        This method moves the object and its underlying graph and graph features to
        the specified device, dtype, or format, and returns the updated object.

        Parameters
        ----------
        *args : Any
            Positional arguments to be passed to the `torch._C._nn._parse_to` function.
        **kwargs : Any
            Keyword arguments to be passed to the `torch._C._nn._parse_to` function.

        Returns
        -------
        GraphCastNet
            The updated object after moving to the specified device, dtype, or format.
        """
        self = super(GraphCastNet, self).to(*args, **kwargs)

        self.g2m_edata = self.g2m_edata.to(*args, **kwargs)
        self.m2g_edata = self.m2g_edata.to(*args, **kwargs)
        self.mesh_ndata = self.mesh_ndata.to(*args, **kwargs)
        self.mesh_edata = self.mesh_edata.to(*args, **kwargs)

        device, _, _, _ = torch._C._nn._parse_to(*args, **kwargs)
        self.g2m_graph = self.g2m_graph.to(device)
        self.mesh_graph = self.mesh_graph.to(device)
        self.m2g_graph = self.m2g_graph.to(device)

        return self