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