PrithviWxC/model.py

from functools import cached_property
from importlib.metadata import version

from torch import Tensor
from torch.utils.checkpoint import checkpoint

if version("torch") > "2.3.0":
    from torch.nn.attention import SDPBackend, sdpa_kernel
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F


# DropPath code is straight from timm
# (https://huggingface.co/spaces/Roll20/pet_score/blame/main/lib/timm/models/layers/drop.py)
def drop_path(
    x: Tensor,
    drop_prob: float = 0.0,
    training: bool = False,
    scale_by_keep: bool = True,
) -> Tensor:
    """Drop paths (Stochastic Depth) per sample (when applied in main path of
    residual blocks). Taken form timm.

    Args:
        x (Tensor): Input tensor.
        drop_prob (float): Probability of dropping `x`, defaults to 0.
        training (bool): Whether model is in in traingin of eval mode,
            defaults to False.
        scale_by_keep (bool): Whether the output should scaled by
            (`1 - drop_prob`), defaults to True.
    Returns:
        Tensor: Tensor that may have randomly dropped with proability
            `drop_path`
    """
    if drop_prob == 0.0 or not training:
        return x
    keep_prob = 1 - drop_prob
    shape = (x.shape[0],) + (1,) * (x.ndim - 1)
    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
    if keep_prob > 0.0 and scale_by_keep:
        random_tensor.div_(keep_prob)
    return x * random_tensor


class DropPath(nn.Module):
    """
    Drop paths (Stochastic Depth) per sample (when applied in main path of
    residual blocks).
    """

    def __init__(
        self, drop_prob: float | None = None, scale_by_keep: bool = True
    ) -> None:
        super(DropPath, self).__init__()
        self.drop_prob = drop_prob
        self.scale_by_keep = scale_by_keep

    def forward(self, x: Tensor) -> Tensor:
        """Runs drop path on input tensor

        Args:
            x: input

        Returns:
            tensor: output after drop_path
        """
        return drop_path(x, self.drop_prob, self.training, self.scale_by_keep)


class Mlp(nn.Module):
    """
    Multi layer perceptron.
    """

    def __init__(
        self, features: int, hidden_features: int, dropout: float = 0.0
    ) -> None:
        """
        Args:
            features: Input/output dimension.
            hidden_features: Hidden dimension.
            dropout: Dropout.
        """
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(features, hidden_features),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_features, features),
            nn.Dropout(dropout),
        )

    def forward(self, x: Tensor) -> Tensor:
        """
        Args:
            x (Tesnor): Tensor of shape [..., channel]
        Returns:
            Tenosr: Tensor of same shape as x.
        """
        return self.net(x)


class LayerNormPassThrough(nn.LayerNorm):
    """Normalising layer that allows the attention mask to be passed through"""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def forward(
        self, d: tuple[Tensor, Tensor | None]
    ) -> tuple[Tensor, Tensor | None]:
        """Forwards function

        Args:
            d (tuple): tuple of the data tensor and the attention mask
        Returns:
            output (Tensor): normalised output data
            attn_mask (Tensor): the attention mask that was passed in
        """
        input, attn_mask = d
        output = F.layer_norm(
            input, self.normalized_shape, self.weight, self.bias, self.eps
        )
        return output, attn_mask


class MultiheadAttention(nn.Module):
    """Multihead attention layer for inputs of shape
    [..., sequence, features].
    """

    def __init__(self, features: int, n_heads: int, dropout: float) -> None:
        """
        Args:
            features: Number of features for inputs to the layer.
            n_heads: Number of attention heads. Should be a factor of features.
                (I.e. the layer uses features // n_heads.)
            dropout: Dropout.
        """  # noqa: E501
        super().__init__()

        if (features % n_heads) != 0:
            raise ValueError(
                f"Features '{features}' is not divisible by heads '{n_heads}'."
            )

        self.features = features
        self.n_heads = n_heads
        self.dropout = dropout

        self.qkv_layer = torch.nn.Linear(features, features * 3, bias=False)
        self.w_layer = torch.nn.Linear(features, features, bias=False)

    def forward(self, d: tuple[Tensor, Tensor | None]) -> Tensor:
        """
        Args:
            d (tuple): tuple containing Tensor of shape [..., sequence, features] and the attention mask
        Returns:
            Tensor: Tensor of shape [..., sequence, features]
        """  # noqa: E501
        x, attn_mask = d

        if not x.shape[-1] == self.features:
            raise ValueError(
                f"Expecting tensor with last dimension size {self.features}."
            )

        passenger_dims = x.shape[:-2]
        B = passenger_dims.numel()
        S = x.shape[-2]
        C = x.shape[-1]
        x = x.reshape(B, S, C)

        # x [B, S, C]
        # q, k, v [B, H, S, C/H]
        q, k, v = (
            self.qkv_layer(x)
            .view(B, S, self.n_heads, 3 * (C // self.n_heads))
            .transpose(1, 2)
            .chunk(chunks=3, dim=3)
        )

        # Let us enforce either flash (A100+) or memory efficient attention.
        if version("torch") > "2.3.0":
            with sdpa_kernel(
                [SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION]
            ):
                # x [B, H, S, C//H]
                x = F.scaled_dot_product_attention(
                    q, k, v, attn_mask=attn_mask, dropout_p=self.dropout
                )
        else:
            with torch.backends.cuda.sdp_kernel(
                enable_flash=True, enable_math=False, enable_mem_efficient=True
            ):
                # x [B, H, S, C//H]
                x = F.scaled_dot_product_attention(
                    q, k, v, dropout_p=self.dropout
                )

        # x [B, S, C]
        x = x.transpose(1, 2).view(B, S, C)

        # x [B, S, C]
        x = self.w_layer(x)

        # Back to input shape
        x = x.view(*passenger_dims, S, self.features)
        return x


class Transformer(nn.Module):
    """
    Transformer for inputs of shape [..., S, features].
    """

    def __init__(
        self,
        features: int,
        mlp_multiplier: int,
        n_heads: int,
        dropout: float,
        drop_path: float,
    ) -> None:
        """
        Args:
            features: Number of features for inputs to the layer.
            mlp_multiplier: Model uses features*mlp_multiplier hidden units.
            n_heads: Number of attention heads. Should be a factor of features.
            (I.e. the layer uses features // n_heads.) dropout: Dropout.
            drop_path: DropPath.
        """
        super().__init__()

        self.features = features
        self.mlp_multiplier = mlp_multiplier
        self.n_heads = n_heads
        self.dropout = dropout
        self.drop_path = (
            DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
        )

        self.attention = nn.Sequential(
            LayerNormPassThrough(features),
            MultiheadAttention(features, n_heads, dropout),
        )

        self.ff = nn.Sequential(
            nn.LayerNorm(features),
            Mlp(
                features=features,
                hidden_features=features * mlp_multiplier,
                dropout=dropout,
            ),
        )

    def forward(self, d: tuple[Tensor, Tensor | None]) -> Tensor:
        """
        Args:
            x: Tensor of shape [..., sequence, features]
        Returns:
            Tensor: Tensor of shape [..., sequence, features]
        """
        x, attn_mask = d
        if not x.shape[-1] == self.features:
            raise ValueError(
                f"Expecting tensor with last dimension size {self.features}."
            )

        attention_x = self.attention(d)

        x = x + self.drop_path(attention_x)
        x = x + self.drop_path(self.ff(x))

        return x


class _Shift(nn.Module):
    """Private base class for the shifter. This allows some behaviour to be
    easily handled when the shifter isn't used.
    """

    def __init__(self):
        super().__init__()

        self._shifted = False

    @torch.no_grad()
    def reset(self) -> None:
        """
        Resets the bool tracking whether the data is shifted
        """
        self._shifted: bool = False

    def forward(self, data: Tensor) -> tuple[Tensor, dict[bool, None]]:
        return data, {True: None, False: None}


class SWINShift(_Shift):
    """
    Handles the shifting of patches similar to how SWIN works. However if we
    shift the latitudes then the poles will wrap and potentially that might be
    problematic. The possition tokens should handle it but masking is safer.
    """

    def __init__(
        self,
        mu_shape: tuple[int, int],
        global_shape: tuple[int, int],
        local_shape: tuple[int, int],
        patch_shape: tuple[int, int],
        n_context_tokens: int = 2,
    ) -> None:
        """
        Args:
            mu_shape: the shape to the masking units
            global_shape: number of global patches in lat and lon
            local_shape: size of the local patches
            patch_shape: patch size
            n_context_token: number of additional context tokens at start of
            _each_ local sequence
        """
        super().__init__()

        self._mu_shape = ms = mu_shape
        self._g_shape = gs = global_shape
        self._l_shape = ls = local_shape
        self._p_shape = ps = patch_shape
        self._lat_patch = (gs[0], ls[0], gs[1], ls[1])
        self._n_context_tokens = n_context_tokens

        self._g_shift_to = tuple(
            int(0.5 * x / p) for x, p in zip(ms, ps, strict=False)
        )
        self._g_shift_from = tuple(
            -int(0.5 * x / p) for x, p in zip(ms, ps, strict=False)
        )

        # Define the attention masks for the shifted MaxViT.
        nglobal = global_shape[0] * global_shape[1]
        nlocal = (
            local_shape[0] * local_shape[1] + self._n_context_tokens
        )  # "+ 1" for leadtime

        lm = torch.ones((nglobal, 1, nlocal, nlocal), dtype=bool)
        mwidth = int(0.5 * local_shape[1]) * local_shape[0]
        lm[
            : gs[1],
            :,
            self._n_context_tokens : mwidth + self._n_context_tokens,
            self._n_context_tokens : mwidth + self._n_context_tokens,
        ] = False
        self.register_buffer("local_mask", lm)

        gm = torch.ones((nlocal, 1, nglobal, nglobal), dtype=bool)
        gm[: int(0.5 * ls[1]) * ls[0], :, : gs[1], : gs[1]] = False
        self.register_buffer("global_mask", gm)

    def _to_grid_global(self, x: Tensor) -> Tensor:
        """
        Shuffle and reshape the data from the global/local setting back to the
        lat/lon grid setting
        Args:
            x: the data tensor to be shuffled.
        Returns:
            x: data in the global/local setting
        """
        nbatch, *other = x.shape

        y1 = x.view(nbatch, *self._g_shape, *self._l_shape, -1)
        y2 = y1.permute(0, 5, 1, 3, 2, 4).contiguous()

        s = y2.shape
        return y2.view((nbatch, -1, s[2] * s[3], s[4] * s[5]))

    def _to_grid_local(self, x: Tensor) -> Tensor:
        """
        Shuffle and reshape the data from the local/global setting to the
        lat/lon grid setting
        Args:
            x: the data tensor to be shuffled.
        Returns:
            x: data in the lat/lon setting.
        """
        x = x.transpose(2, 1).contiguous()
        return self._to_grid_global(x)

    def _from_grid_global(self, x: Tensor) -> Tensor:
        """
        Shuffle and reshape the data from the lat/lon grid to the global/local
        setting
        Args:
            x: the data tensor to be shuffled.
        Returns:
            x: data in the global/local setting
        """
        nbatch, *other = x.shape

        z1 = x.view(nbatch, -1, *self._lat_patch)
        z2 = z1.permute(0, 2, 4, 3, 5, 1).contiguous()

        s = z2.shape
        return z2.view(nbatch, s[1] * s[2], s[3] * s[4], -1)

    def _from_grid_local(self, x: Tensor) -> Tensor:
        """
        Shuffle and reshape the data from the lat/lon grid to the local/global
        setting
        Args:
            x: the data tensor to be shuffled.
        Returns:
            x: data in the local/global setting
        """
        x = self._from_grid_global(x)
        return x.transpose(2, 1).contiguous()

    def _shift(self, x: Tensor) -> Tensor:
        """
        Shifts data in the gridded lat/lon setting by half the mask unit shape
        Args:
            x: data to be shifted
        Returns:
            x: either the hsifted or unshifted data
        """
        shift = self._g_shift_from if self._shifted else self._g_shift_to
        x_shifted = torch.roll(x, shift, (-2, -1))

        self._shifted = not self._shifted
        return x_shifted

    def _sep_lt(self, x: Tensor) -> tuple[Tensor, Tensor]:
        """
        Seperate off the leadtime from the local patches
        Args:
            x: data to have leadtime removed from
        Returns:
            lt: leadtime
            x: data without the lead time in the local patch
        """
        lt_it = x[:, : self._n_context_tokens, :, :]
        x_stripped = x[:, self._n_context_tokens :, :, :]

        return lt_it, x_stripped

    def forward(self, data: Tensor) -> tuple[Tensor, Tensor]:
        """Shift or unshift the the data depending on whether the data is
        already shifted, as defined by self._shifte.

        Args:
            data: data to be shifted
        Returns:
            Tensor: shifted data Tensor
        """
        lt, x = self._sep_lt(data)

        x_grid = self._to_grid_local(x)
        x_shifted = self._shift(x_grid)
        x_patched = self._from_grid_local(x_shifted)

        # Mask has to be repeated based on batch size
        n_batch = x_grid.shape[0]
        local_rep = [n_batch] + [1] * (self.local_mask.ndim - 1)
        global_rep = [n_batch] + [1] * (self.global_mask.ndim - 1)

        if self._shifted:
            attn_mask = {
                True: self.local_mask.repeat(local_rep),
                False: self.global_mask.repeat(global_rep),
            }
        else:
            attn_mask = {True: None, False: None}

        return torch.cat((lt, x_patched), axis=1), attn_mask


class LocalGlobalLocalBlock(nn.Module):
    """
    Applies alternating block and grid attention. Given a parameter n_blocks,
    the entire module contains 2*n_blocks+1 transformer blocks. The first,
    third, ..., last apply local (block) attention. The second, fourth, ...
    global (grid) attention.

    This is heavily inspired by
    Tu et al. "MaxViT: Multi-Axis Vision Transformer"
    (https://arxiv.org/abs/2204.01697).
    """

    def __init__(
        self,
        features: int,
        mlp_multiplier: int,
        n_heads: int,
        dropout: float,
        n_blocks: int,
        drop_path: float,
        shifter: nn.Module | None = None,
        checkpoint: list[int] | None = None,
    ) -> None:
        """
        Args:
            features: Number of features for inputs to the layer.
            mlp_multiplier: Model uses features*mlp_multiplier hidden units.
            n_heads: Number of attention heads. Should be a factor of features.
            (I.e. the layer uses features // n_heads.)
            dropout: Dropout.
            drop_path: DropPath.
            n_blocks: Number of local-global transformer pairs.
        """
        super().__init__()

        self.features = features
        self.mlp_multiplier = mlp_multiplier
        self.n_heads = n_heads
        self.dropout = dropout
        self.drop_path = drop_path
        self.n_blocks = n_blocks
        self._checkpoint = checkpoint or []

        if not all(0 <= c < 2 * n_blocks + 1 for c in self._checkpoint):
            raise ValueError(
                "Checkpoints should be 0 <= i < 2*n_blocks+1. "
                f"{self._checkpoint=}."
            )

        self.transformers = nn.ModuleList(
            [
                Transformer(
                    features=features,
                    mlp_multiplier=mlp_multiplier,
                    n_heads=n_heads,
                    dropout=dropout,
                    drop_path=drop_path,
                )
                for _ in range(2 * n_blocks + 1)
            ]
        )

        self.evaluator = [
            self._checkpoint_wrapper
            if i in self._checkpoint
            else lambda m, x: m(x)
            for i, _ in enumerate(self.transformers)
        ]

        self.shifter = shifter or _Shift()

    @staticmethod
    def _checkpoint_wrapper(
        model: nn.Module, data: tuple[Tensor, Tensor | None]
    ) -> Tensor:
        return checkpoint(model, data, use_reentrant=False)

    def forward(self, x: Tensor) -> Tensor:
        """
        Args:
            x: Tensor of shape::

                [batch, global_sequence, local_sequence, features]

        Returns:
            Tensor: Tensor of shape::

                [batch, global_sequence, local_sequence, features]
        """
        if x.shape[-1] != self.features:
            raise ValueError(
                f"Expecting tensor with last dimension size {self.features}."
            )
        if x.ndim != 4:
            raise ValueError(
                f"Expecting tensor with exactly four dimensions. {x.shape=}."
            )

        self.shifter.reset()
        local: bool = True
        attn_mask = {True: None, False: None}

        transformer_iter = zip(self.evaluator, self.transformers, strict=False)

        # First local block
        evaluator, transformer = next(transformer_iter)
        x = evaluator(transformer, (x, attn_mask[local]))

        for evaluator, transformer in transformer_iter:
            local = not local
            # We are making exactly 2*n_blocks transposes.
            # So the output has the same shape as input.
            x = x.transpose(1, 2)

            x = evaluator(transformer, (x, attn_mask[local]))

            if not local:
                x, attn_mask = self.shifter(x)

        return x


class PatchEmbed(nn.Module):
    """
    Patch embedding via 2D convolution.
    """

    def __init__(
        self, patch_size: int | tuple[int, ...], channels: int, embed_dim: int
    ):
        super().__init__()

        self.patch_size = patch_size
        self.channels = channels
        self.embed_dim = embed_dim

        self.proj = nn.Conv2d(
            channels,
            embed_dim,
            kernel_size=patch_size,
            stride=patch_size,
            bias=True,
        )

    def forward(self, x: Tensor) -> Tensor:
        """
        Args:
            x: Tensor of shape [batch, channels, lat, lon].
        Returns:
            Tensor: Tensor with shape
                [batch, embed_dim, lat//patch_size, lon//patch_size]
        """

        H, W = x.shape[-2:]

        if W % self.patch_size[1] != 0:
            raise ValueError(
                f"Cannot do patch embedding for tensor of shape {x.size()}"
                " with patch size {self.patch_size}. (Dimensions are BSCHW.)"
            )
        if H % self.patch_size[0] != 0:
            raise ValueError(
                f"Cannot do patch embedding for tensor of shape {x.size()}"
                f" with patch size {self.patch_size}. (Dimensions are BSCHW.)"
            )

        x = self.proj(x)

        return x


class PrithviWxCEncoderDecoder(nn.Module):
    """
    Hiera-MaxViT encoder/decoder code.
    """

    def __init__(
        self,
        embed_dim: int,
        n_blocks: int,
        mlp_multiplier: float,
        n_heads: int,
        dropout: float,
        drop_path: float,
        shifter: nn.Module | None = None,
        transformer_cp: list[int] | None = None,
    ) -> None:
        """
        Args:
            embed_dim: Embedding dimension
            n_blocks: Number of local-global transformer pairs.
            mlp_multiplier: MLP multiplier for hidden features in feed forward
                networks.
            n_heads: Number of attention heads.
            dropout: Dropout.
            drop_path: DropPath.
        """
        super().__init__()

        self.embed_dim = embed_dim
        self.n_blocks = n_blocks
        self.mlp_multiplier = mlp_multiplier
        self.n_heads = n_heads
        self.dropout = dropout
        self._transformer_cp = transformer_cp

        self.lgl_block = LocalGlobalLocalBlock(
            features=embed_dim,
            mlp_multiplier=mlp_multiplier,
            n_heads=n_heads,
            dropout=dropout,
            drop_path=drop_path,
            n_blocks=n_blocks,
            shifter=shifter,
            checkpoint=transformer_cp,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: Tensor of shape
              [batch, global sequence, local sequence, embed_dim]
        Returns:
            Tensor of shape
                [batch, mask_unit_sequence, local_sequence, embed_dim].
                Identical in shape to the input x.
        """

        x = self.lgl_block(x)

        return x


class PrithviWxC(nn.Module):
    """Encoder-decoder fusing Hiera with MaxViT. See
    - Ryali et al. "Hiera: A Hierarchical Vision Transformer without the
    Bells-and-Whistles" (https://arxiv.org/abs/2306.00989)
    - Tu et al. "MaxViT: Multi-Axis Vision Transformer"
    (https://arxiv.org/abs/2204.01697)
    """

    def __init__(
        self,
        in_channels: int,
        input_size_time: int,
        in_channels_static: int,
        input_scalers_mu: Tensor,
        input_scalers_sigma: Tensor,
        input_scalers_epsilon: float,
        static_input_scalers_mu: Tensor,
        static_input_scalers_sigma: Tensor,
        static_input_scalers_epsilon: float,
        output_scalers: Tensor,
        n_lats_px: int,
        n_lons_px: int,
        patch_size_px: tuple[int],
        mask_unit_size_px: tuple[int],
        mask_ratio_inputs: float,
        embed_dim: int,
        n_blocks_encoder: int,
        n_blocks_decoder: int,
        mlp_multiplier: float,
        n_heads: int,
        dropout: float,
        drop_path: float,
        parameter_dropout: float,
        residual: str,
        masking_mode: str,
        positional_encoding: str,
        decoder_shifting: bool = False,
        checkpoint_encoder: list[int] | None = None,
        checkpoint_decoder: list[int] | None = None,
    ) -> None:
        """
        Args:
            in_channels: Number of input channels.
            input_size_time: Number of timestamps in input.
            in_channels_static: Number of input channels for static data.
            input_scalers_mu: Tensor of size (in_channels,). Used to rescale
                input.
            input_scalers_sigma: Tensor of size (in_channels,). Used to rescale
                input.
            input_scalers_epsilon: Float. Used to rescale input.
            static_input_scalers_mu: Tensor of size (in_channels_static). Used
                to rescale static inputs.
            static_input_scalers_sigma: Tensor of size (in_channels_static).
                Used to rescale static inputs.
            static_input_scalers_epsilon: Float. Used to rescale static inputs.
            output_scalers: Tensor of shape (in_channels,). Used to rescale
                output.
            n_lats_px: Total latitudes in data. In pixels.
            n_lons_px: Total longitudes in data. In pixels.
            patch_size_px: Patch size for tokenization. In pixels lat/lon.
            mask_unit_size_px: Size of each mask unit. In pixels lat/lon.
            mask_ratio_inputs: Masking ratio for inputs. 0 to 1.
            embed_dim: Embedding dimension
            n_blocks_encoder: Number of local-global transformer pairs in
                encoder.
            n_blocks_decoder: Number of local-global transformer pairs in
                decoder.
            mlp_multiplier: MLP multiplier for hidden features in feed forward
                networks.
            n_heads: Number of attention heads.
            dropout: Dropout.
            drop_path: DropPath.
            parameter_dropout: Dropout applied to parameters.
            residual: Indicates whether and how model should work as residual
                model. Accepted values are 'climate', 'temporal' and 'none'
            positional_encoding: possible values are
              ['absolute' (default), 'fourier'].
                'absolute'  lat lon encoded in 3 dimensions using sine and
                  cosine
                'fourier' lat/lon to be encoded using various frequencies
            masking_mode: String ['local', 'global', 'both'] that controls the
                type of masking used.
            checkpoint_encoder: List of integers controlling if gradient
              checkpointing is used on encoder.
                Format: [] for no gradient checkpointing. [3, 7] for
                  checkpointing after 4th and 8th layer etc.
            checkpoint_decoder: List of integers controlling if gradient
              checkpointing is used on decoder.
                Format: See `checkpoint_encoder`.
            masking_mode: The type of masking to use
              {'global', 'local', 'both'}
            decoder_shifting: Whether to use swin shifting in the decoder.
        """
        super().__init__()

        self.in_channels = in_channels
        self.input_size_time = input_size_time
        self.in_channels_static = in_channels_static
        self.n_lats_px = n_lats_px
        self.n_lons_px = n_lons_px
        self.patch_size_px = patch_size_px
        self.mask_unit_size_px = mask_unit_size_px
        self.mask_ratio_inputs = mask_ratio_inputs
        self.embed_dim = embed_dim
        self.n_blocks_encoder = n_blocks_encoder
        self.n_blocks_decoder = n_blocks_decoder
        self.mlp_multiplier = mlp_multiplier
        self.n_heads = n_heads
        self.dropout = dropout
        self.drop_path = drop_path
        self.residual = residual
        self._decoder_shift = decoder_shifting
        self.positional_encoding = positional_encoding
        self._checkpoint_encoder = checkpoint_encoder
        self._checkpoint_decoder = checkpoint_decoder

        assert self.n_lats_px % self.mask_unit_size_px[0] == 0
        assert self.n_lons_px % self.mask_unit_size_px[1] == 0
        assert self.mask_unit_size_px[0] % self.patch_size_px[0] == 0
        assert self.mask_unit_size_px[1] % self.patch_size_px[1] == 0

        if self.patch_size_px[0] != self.patch_size_px[1]:
            raise NotImplementedError(
                "Current pixel shuffle symmetric patches."
            )

        self.local_shape_mu = (
            self.mask_unit_size_px[0] // self.patch_size_px[0],
            self.mask_unit_size_px[1] // self.patch_size_px[1],
        )
        self.global_shape_mu = (
            self.n_lats_px // self.mask_unit_size_px[0],
            self.n_lons_px // self.mask_unit_size_px[1],
        )

        assert input_scalers_mu.shape == (in_channels,)
        assert input_scalers_sigma.shape == (in_channels,)
        assert output_scalers.shape == (in_channels,)

        if self.positional_encoding != "fourier":
            assert static_input_scalers_mu.shape == (in_channels_static,)
            assert static_input_scalers_sigma.shape == (in_channels_static,)

        # Input shape [batch, time, parameter, lat, lon]
        self.input_scalers_epsilon = input_scalers_epsilon
        self.register_buffer(
            "input_scalers_mu", input_scalers_mu.reshape(1, 1, -1, 1, 1)
        )
        self.register_buffer(
            "input_scalers_sigma", input_scalers_sigma.reshape(1, 1, -1, 1, 1)
        )

        # Static inputs shape [batch, parameter, lat, lon]
        self.static_input_scalers_epsilon = static_input_scalers_epsilon
        self.register_buffer(
            "static_input_scalers_mu",
            static_input_scalers_mu.reshape(1, -1, 1, 1),
        )
        self.register_buffer(
            "static_input_scalers_sigma",
            static_input_scalers_sigma.reshape(1, -1, 1, 1),
        )

        # Output shape [batch, parameter, lat, lon]
        self.register_buffer(
            "output_scalers", output_scalers.reshape(1, -1, 1, 1)
        )

        self.parameter_dropout = nn.Dropout2d(p=parameter_dropout)

        self.patch_embedding = PatchEmbed(
            patch_size=patch_size_px,
            channels=in_channels * input_size_time,
            embed_dim=embed_dim,
        )

        if self.residual == "climate":
            self.patch_embedding_static = PatchEmbed(
                patch_size=patch_size_px,
                channels=in_channels + in_channels_static,
                embed_dim=embed_dim,
            )
        else:
            self.patch_embedding_static = PatchEmbed(
                patch_size=patch_size_px,
                channels=in_channels_static,
                embed_dim=embed_dim,
            )

        self.input_time_embedding = nn.Linear(1, embed_dim // 4, bias=True)
        self.lead_time_embedding = nn.Linear(1, embed_dim // 4, bias=True)

        self.mask_token = nn.Parameter(torch.randn(1, 1, 1, self.embed_dim))
        self._nglobal_mu = np.prod(self.global_shape_mu)
        self._global_idx = torch.arange(self._nglobal_mu)

        self._nlocal_mu = np.prod(self.local_shape_mu)
        self._local_idx = torch.arange(self._nlocal_mu)

        self.encoder = PrithviWxCEncoderDecoder(
            embed_dim=embed_dim,
            n_blocks=n_blocks_encoder,
            mlp_multiplier=mlp_multiplier,
            n_heads=n_heads,
            dropout=dropout,
            drop_path=drop_path,
            transformer_cp=checkpoint_encoder,
        )

        if n_blocks_decoder != 0:
            if self._decoder_shift:
                self.decoder_shifter = d_shifter = SWINShift(
                    self.mask_unit_size_px,
                    self.global_shape_mu,
                    self.local_shape_mu,
                    self.patch_size_px,
                    n_context_tokens=0,
                )
            else:
                self.decoder_shifter = d_shifter = None

            self.decoder = PrithviWxCEncoderDecoder(
                embed_dim=embed_dim,
                n_blocks=n_blocks_decoder,
                mlp_multiplier=mlp_multiplier,
                n_heads=n_heads,
                dropout=dropout,
                drop_path=0.0,
                shifter=d_shifter,
                transformer_cp=checkpoint_decoder,
            )

            self.unembed = nn.Linear(
                self.embed_dim,
                self.in_channels
                * self.patch_size_px[0]
                * self.patch_size_px[1],
                bias=True,
            )

        self.masking_mode = masking_mode.lower()
        match self.masking_mode:
            case "local":
                self.generate_mask = self._gen_mask_local
            case "global":
                self.generate_mask = self._gen_mask_global
            case "both":
                self._mask_both_local: bool = True
                self.generate_mask = self._gen_mask_both
            case _:
                raise ValueError(
                    f"Masking mode '{masking_mode}' not supported"
                )

    def swap_masking(self) -> None:
        self._mask_both_local = not self._mask_both_local

    @cached_property
    def n_masked_global(self):
        return int(self.mask_ratio_inputs * np.prod(self.global_shape_mu))

    @cached_property
    def n_masked_local(self):
        return int(self.mask_ratio_inputs * np.prod(self.local_shape_mu))

    @staticmethod
    def _shuffle_along_axis(a, axis):
        idx = torch.argsort(input=torch.rand(*a.shape), dim=axis)
        return torch.gather(a, dim=axis, index=idx)

    def _gen_mask_local(self, sizes: tuple[int]) -> tuple[Tensor]:
        """
        Args:
            batch_size: Number of elements in batch
        Returns:
            Tuple of torch tensors. [indices masked, indices unmasked].
            Each of these is a tensor of shape (batch, global sequene)
        """
        # Identify which indices (values) should be masked

        maskable_indices = self._local_idx.view(1, -1).expand(*sizes[:2], -1)

        maskable_indices = self._shuffle_along_axis(maskable_indices, 2)

        indices_masked = maskable_indices[:, :, : self.n_masked_local]
        indices_unmasked = maskable_indices[:, :, self.n_masked_local :]

        return indices_masked, indices_unmasked

    def _gen_mask_global(self, sizes: tuple[int]) -> tuple[Tensor]:
        """
        Args:
            batch_size: Number of elements in batch
        Returns:
            Tuple of torch tensors. [indices masked, indices unmasked].
            Each of these is a tensor of shape (batch, global sequene)
        """
        # Identify which indices (values) should be masked

        maskable_indices = self._global_idx.view(1, -1).expand(*sizes[:1], -1)

        maskable_indices = self._shuffle_along_axis(maskable_indices, 1)

        indices_masked = maskable_indices[:, : self.n_masked_global]
        indices_unmasked = maskable_indices[:, self.n_masked_global :]

        return indices_masked, indices_unmasked

    def _gen_mask_both(self, sizes: tuple[int]) -> tuple[Tensor]:
        if self._mask_both_local:
            return self._gen_mask_local(sizes)
        else:
            return self._gen_mask_global(sizes)

    @staticmethod
    def reconstruct_batch(
        idx_masked: Tensor,
        idx_unmasked: Tensor,
        data_masked: Tensor,
        data_unmasked: Tensor,
    ) -> Tensor:
        """Reconstructs a tensor along the mask unit dimension. Batched
        version.

        Args:
            idx_masked: Tensor of shape `batch, mask unit sequence`.
            idx_unmasked: Tensor of shape `batch, mask unit sequence`.
            data_masked: Tensor of shape `batch, mask unit sequence, ...`.
                Should have same size along mask unit sequence dimension as
                idx_masked. Dimensions beyond the first two, marked here as ...
                will typically be `local_sequence, channel` or
                `channel, lat, lon`. These dimensions should agree with
                data_unmasked.
            data_unmasked: Tensor of shape `batch, mask unit sequence, ...`.
                Should have same size along mask unit sequence dimension as
                idx_unmasked. Dimensions beyond the first two, marked here as
                ... will typically be `local_sequence, channel` or `channel,
                lat, lon`. These dimensions should agree with data_masked.
        Returns:
            Tensor: Tensor of same shape as inputs data_masked and
                data_unmasked. I.e. `batch, mask unit sequence, ...`. Index for
                the total data composed of the masked and the unmasked part.
        """
        dim: int = idx_masked.ndim

        idx_total = torch.argsort(
            torch.cat([idx_masked, idx_unmasked], dim=-1), dim=-1
        )
        idx_total = idx_total.view(
            *idx_total.shape, *[1] * (data_unmasked.ndim - dim)
        )
        idx_total = idx_total.expand(
            *idx_total.shape[:dim], *data_unmasked.shape[dim:]
        )

        data = torch.cat([data_masked, data_unmasked], dim=dim - 1)
        data = torch.gather(data, dim=dim - 1, index=idx_total)

        return data, idx_total

    def fourier_pos_encoding(self, x_static: Tensor) -> Tensor:
        """
        Args
            x_static: B x C x H x W. first two channels are lat, and lon
        Returns
            Tensor: Tensor of shape B x E x H x W where E is the embedding
                dimension.
        """

        # B x C x H x W -> B x 1 x H/P x W/P
        latitudes_patch = F.avg_pool2d(
            x_static[:, [0]],
            kernel_size=self.patch_size_px,
            stride=self.patch_size_px,
        )
        longitudes_patch = F.avg_pool2d(
            x_static[:, [1]],
            kernel_size=self.patch_size_px,
            stride=self.patch_size_px,
        )

        modes = (
            torch.arange(self.embed_dim // 4, device=x_static.device).view(
                1, -1, 1, 1
            )
            + 1.0
        )
        pos_encoding = torch.cat(
            (
                torch.sin(latitudes_patch * modes),
                torch.sin(longitudes_patch * modes),
                torch.cos(latitudes_patch * modes),
                torch.cos(longitudes_patch * modes),
            ),
            axis=1,
        )

        return pos_encoding  # B x E x H/P x W/P

    def time_encoding(self, input_time, lead_time):
        """
        Args:
            input_time: Tensor of shape [batch].
            lead_time: Tensor of shape [batch].
        Returns:
            Tensor: Tensor of shape [batch, embed_dim, 1, 1]
        """
        input_time = self.input_time_embedding(input_time.view(-1, 1, 1, 1))
        lead_time = self.lead_time_embedding(lead_time.view(-1, 1, 1, 1))

        time_encoding = torch.cat(
            (
                torch.cos(input_time),
                torch.cos(lead_time),
                torch.sin(input_time),
                torch.sin(lead_time),
            ),
            axis=3,
        )
        return time_encoding

    def to_patching(self, x: Tensor) -> Tensor:
        """Transform data from lat/lon space to two axis patching

        Args: ->
            x: Tesnor in lat/lon space (N, C, Nlat//P_0, Nlon//P_1)

        Returns:
            Tensor in patch space (N, G, L, C)
        """
        n_batch = x.shape[0]

        x = x.view(
            n_batch,
            -1,
            self.global_shape_mu[0],
            self.local_shape_mu[0],
            self.global_shape_mu[1],
            self.local_shape_mu[1],
        )
        x = x.permute(0, 2, 4, 3, 5, 1).contiguous()

        s = x.shape
        return x.view(n_batch, s[1] * s[2], s[3] * s[4], -1)

    def from_patching(self, x: Tensor) -> Tensor:
        """Transform data from two axis patching to lat/lon space

        Args:
            x: Tensor in patch space with shape (N, G, L, C*P_0*P_1)

        Returns:
            Tensor: Tensor in lat/lon space
                (N, C*P_0*P_1, Nlat//P_0, Nlon // P_1)
        """
        n_batch = x.shape[0]

        x = x.view(
            n_batch,
            self.global_shape_mu[0],
            self.global_shape_mu[1],
            self.local_shape_mu[0],
            self.local_shape_mu[1],
            -1,
        )
        x = x.permute(0, 5, 1, 3, 2, 4).contiguous()

        s = x.shape
        return x.view(n_batch, -1, s[2] * s[3], s[4] * s[5])

    def forward(self, batch: dict[str, torch.Tensor]) -> torch.Tensor:
        """
        Args:
            batch: Dictionary the following keys::

                'x': Tensor of shape [batch, time, parameter, lat, lon]
                'y': Tensor of shape [batch, parameter, lat, lon]
                'static': Tensor of shape [batch, channel_static, lat, lon]
                'climate': Optional tensor of shape [batch, parameter, lat, lon]
                'input_time': Tensor of shape [batch]. Or none.
                'lead_time': Tensor of shape [batch]. Or none.

        Returns:
            Tensor: Tensor of shape [batch, parameter, lat, lon].
        """  # noqa: E501
        x_rescaled = (batch["x"] - self.input_scalers_mu) / (
            self.input_scalers_sigma + self.input_scalers_epsilon
        )
        batch_size = x_rescaled.shape[0]

        if self.positional_encoding == "fourier":
            x_static_pos = self.fourier_pos_encoding(batch["static"])
            x_static = (
                batch["static"][:, 2:] - self.static_input_scalers_mu[:, 3:]
            ) / (
                self.static_input_scalers_sigma[:, 3:]
                + self.static_input_scalers_epsilon
            )
        else:
            x_static = (batch["static"] - self.static_input_scalers_mu) / (
                self.static_input_scalers_sigma
                + self.static_input_scalers_epsilon
            )

        if self.residual == "temporal":
            # We create a residual of same shape as y
            index = torch.where(
                batch["lead_time"] > 0, batch["x"].shape[1] - 1, 0
            )
            index = index.view(-1, 1, 1, 1, 1)
            index = index.expand(batch_size, 1, *batch["x"].shape[2:])
            x_hat = torch.gather(batch["x"], dim=1, index=index)
            x_hat = x_hat.squeeze(1)
        elif self.residual == "climate":
            climate_scaled = (
                batch["climate"] - self.input_scalers_mu.view(1, -1, 1, 1)
            ) / (
                self.input_scalers_sigma.view(1, -1, 1, 1)
                + self.input_scalers_epsilon
            )

        # [batch, time, parameter, lat, lon]
        # -> [batch, time x parameter, lat, lon]
        x_rescaled = x_rescaled.flatten(1, 2)
        # Parameter dropout
        x_rescaled = self.parameter_dropout(x_rescaled)

        x_embedded = self.patch_embedding(x_rescaled)

        if self.residual == "climate":
            static_embedded = self.patch_embedding_static(
                torch.cat((x_static, climate_scaled), dim=1)
            )
        else:
            static_embedded = self.patch_embedding_static(x_static)

        if self.positional_encoding == "fourier":
            static_embedded += x_static_pos

        x_embedded = self.to_patching(x_embedded)
        static_embedded = self.to_patching(static_embedded)

        time_encoding = self.time_encoding(
            batch["input_time"], batch["lead_time"]
        )

        tokens = x_embedded + static_embedded + time_encoding

        # Now we generate masks based on masking_mode
        indices_masked, indices_unmasked = self.generate_mask(
            (batch_size, self._nglobal_mu)
        )
        indices_masked = indices_masked.to(device=tokens.device)
        indices_unmasked = indices_unmasked.to(device=tokens.device)
        maskdim: int = indices_masked.ndim

        # Unmasking
        unmask_view = (*indices_unmasked.shape, *[1] * (tokens.ndim - maskdim))
        unmasked = torch.gather(
            tokens,
            dim=maskdim - 1,
            index=indices_unmasked.view(*unmask_view).expand(
                *indices_unmasked.shape, *tokens.shape[maskdim:]
            ),
        )

        # Encoder
        x_encoded = self.encoder(unmasked)

        # Generate and position encode the mask tokens
        # [1, 1, 1, embed_dim]
        # -> [batch, global_seq_masked, local seq, embed_dim]
        mask_view = (*indices_masked.shape, *[1] * (tokens.ndim - maskdim))
        masking = self.mask_token.repeat(*static_embedded.shape[:3], 1)
        masked = masking + static_embedded
        masked = torch.gather(
            masked,
            dim=maskdim - 1,
            index=indices_masked.view(*mask_view).expand(
                *indices_masked.shape, *tokens.shape[maskdim:]
            ),
        )

        recon, _ = self.reconstruct_batch(
            indices_masked, indices_unmasked, masked, x_encoded
        )

        x_decoded = self.decoder(recon)

        # Output: [batch, global sequence, local sequence,
        #          in_channels * patch_size[0] * patch_size[1]]
        x_unembed = self.unembed(x_decoded)

        # Reshape to [batch, global_lat, global_lon, local_lat, local_lon,
        #             in_channels * patch_size[0] * patch_size[1]]
        x_out = self.from_patching(x_unembed)

        # Pixel shuffle to [batch, in_channels, lat, lon]
        x_out = F.pixel_shuffle(x_out, self.patch_size_px[0])

        if self.residual == "temporal":
            x_out = self.output_scalers * x_out + x_hat
        elif self.residual == "climate":
            x_out = self.output_scalers * x_out + batch["climate"]
        elif self.residual == "none":
            x_out = (
                self.output_scalers * x_out
                + self.input_scalers_mu.reshape(1, -1, 1, 1)
            )

        return x_out