model/prithvi_wxc_official.py

from importlib.metadata import version
TORCH_VERSION = version('torch')

from functools import cached_property
from typing import Optional
import numpy as np

import torch
from torch import Tensor
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint
if TORCH_VERSION > '2.3.0':
    from torch.nn.attention import SDPBackend, sdpa_kernel


# DropPath code is straight from timm
# (https://huggingface.co/spaces/Roll20/pet_score/blame/main/lib/timm/models/layers/drop.py)
# Primarily since we currently don't have timm in the environment.
def drop_path(
    x, drop_prob: float = 0.0, training: bool = False, scale_by_keep: bool = True
):
    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).

    This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
    the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
    changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
    'survival rate' as the argument.
    """
    if drop_prob == 0.0 or not training:
        return x
    keep_prob = 1 - drop_prob
    shape = (x.shape[0],) + (1,) * (
        x.ndim - 1
    )  # work with diff dim tensors, not just 2D ConvNets
    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=None, scale_by_keep=True):
        super(DropPath, self).__init__()
        self.drop_prob = drop_prob
        self.scale_by_keep = scale_by_keep

    def forward(self, x):
        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:
            Tensor of shape [..., channel]
        Returns:
            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 of the data tensor and the attention mask
        Returns:
            output: normalised output data
            attn_mask: 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].

    Uses `scaled_dot_product_attention` to obtain a memory efficient attention
    computation (https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html).
    This follows:
    - Dao et la. "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness"
        (https://arxiv.org/abs/2205.14135)
    - Rabe, Staats "Self-attention Does Not Need O(n2) Memory" (https://arxiv.org/abs/2112.05682)

    Note: Even though the documentation page for `scaled_dot_product_attention`
    states that tensors can have any number of dimensions as long as the shapes
    are along the lines of `(B, ..., S, E)`, the fused and memory efficient
    mechanisms we enforce here require a 4D input. Some experimentatino shows
    that this should be of shape `(B, H, S, E)`, where `H` represents heads.
    However, as of right now this is not confirmed int he documentation.
    """

    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.
        """
        super().__init__()

        if not (features % n_heads) == 0:
            raise ValueError(
                f"Number of features {features} is not divisible by number of 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 containing Tensor of shape [..., sequence, features] and
            the attention mask
        Returns:
            Tensor of shape [..., sequence, features]
        """
        x, attn_mask = d

        if not x.shape[-1] == self.features:
            raise ValueError(
                f"Expecting tensor with last dimension of 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,
        # falling back to math on CPU / DCU(昇腾) shim where fused kernels are unavailable.
        if TORCH_VERSION > '2.3.0':
            with sdpa_kernel([SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH]):
                # x [B, H, S, C//H]
                x = F.scaled_dot_product_attention(
                    q, k, v, 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).reshape(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 will use 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 of shape [..., sequence, features]
        """
        x, attn_mask = d
        if not x.shape[-1] == self.features:
            raise ValueError(
                f"Expecting tensor with last dimension of 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))
        self._g_shift_from = tuple(-int(0.5 * x / p) for x, p in zip(ms, ps))

        # 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._shifted
        Args:
            data: data to be shifted
        Returns:

        """
        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 SWINShiftNoBuffer(_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))
        self._g_shift_from = tuple(-int(0.5 * x / p) for x, p in zip(ms, ps))

        # 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.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.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._shifted
        Args:
            data: data to be shifted
        Returns:

        """
        lt, x = self._sep_lt(data)

        if self.local_mask.device != x.device:
            self.local_mask = self.local_mask.to(device=x.device)
        if self.global_mask.device != x.device:
            self.global_mask = self.global_mask.to(device=x.device)

        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:
        """
        Args:
            features: Number of features for inputs to the layer.
            mlp_multiplier: Model will use 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

        if len(checkpoint) > 0:
            if min(checkpoint) < 0 or max(checkpoint) >= 2 * n_blocks + 1:
                raise ValueError(f'Checkpoints should satisfy 0 <= i < 2*n_blocks+1. We have {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 checkpoint else lambda m, x : m(x)
            for i, _ in enumerate(self.transformers)
        ]

        self.shifter = shifter or _Shift()

    @staticmethod
    def _checkpoint_wrapper(model, data):
        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 of shape [batch, global_sequence, local_sequence, features]
        """
        if x.shape[-1] != self.features:
            raise ValueError(
                f"Expecting tensor with last dimension of size {self.features}."
            )
        if x.ndim != 4:
            raise ValueError(
                f"Expecting tensor with exactly four dimensions. Input has shape {x.shape}."
            )

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

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

        # 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 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:
        """
        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,
        mask_ratio_targets: 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,
        encoder_shifting: bool = False,
        decoder_shifting: bool = False,
        checkpoint_encoder: list[int]=[],
        checkpoint_decoder: list[int]=[],
    ) -> 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.
            mask_ratio_targets: Masking ratio for targets. 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'}
            encoder_shifting: Whether to use swin shifting in the encoder.
            decoder_shifting: Whether to use swin shifting in the decoder.
        """
        super().__init__()

        if mask_ratio_targets > 0.0:
            raise NotImplementedError("Target masking is not implemented.")

        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.mask_ratio_targets = mask_ratio_targets
        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._encoder_shift = encoder_shifting
        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 implementation assumes same patch size along both dimensions."
            )

        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)

        if self._encoder_shift:
            self.encoder_shifter = e_shifter = SWINShiftNoBuffer(
                self.mask_unit_size_px,
                self.global_shape_mu,
                self.local_shape_mu,
                self.patch_size_px,
                n_context_tokens=0,
            )
        else:
            self.encoder_shifter = e_shifter = None
        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,
            shifter=e_shifter,
            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.,
                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:
        if hasattr(self, '_mask_both_local'):
            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):
        # https://stackoverflow.com/questions/5040797/shuffling-numpy-array-along-a-given-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)
        """
        # We identifies 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)

        # `...` cannot be jit'd :-(
        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)
        """
        # We identifies 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 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):
        """
        Args
            x_static: B x C x H x W. first two channels are lat, and lon respectively
        Returns
            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.
        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 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,
            self.embed_dim,
            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 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 containing the keys 'x', 'y', 'input_time',
                'lead_time' and 'static'. The associated torch tensors have the
                following shapes:
                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 of shape [batch, parameter, lat, lon].
        """
        assert batch["x"].shape[2] == self.in_channels
        assert batch["x"].shape[3] == self.n_lats_px
        assert batch["x"].shape[4] == self.n_lons_px
        assert batch["y"].shape[1] == self.in_channels
        assert batch["y"].shape[2] == self.n_lats_px
        assert batch["y"].shape[3] == self.n_lons_px
        if self.positional_encoding == 'fourier':
            # the first two features (lat, lon) are encoded separately
            assert batch['static'].shape[1] - 2 == self.in_channels_static, "When setting self.positional_encoding to fourier, the number of static params change in the dataset. So, in the config, reduce num_static_channels (e.g., 4 instead of 7)."
        else:
            assert batch['static'].shape[1] == self.in_channels_static
        assert batch["static"].shape[2] == self.n_lats_px
        assert batch["static"].shape[3] == self.n_lons_px

        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']) # B, embed_dim, lat / patch_size, lon / patch_size
            x_static = (batch['static'][:, 2:] - self.static_input_scalers_mu[:, 3:]) / ( # The first two channels in batch['static'] are used in positional encoding
                self.static_input_scalers_sigma[:, 3:] + self.static_input_scalers_epsilon # This translates to the first three channels in 'static_input_scalers_mu'
            )
        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)
            assert (
                batch["y"].shape == x_hat.shape
            ), f'Shapes {batch["y"].shape} and {x_hat.shape} do not agree.'
        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)
        assert x_embedded.shape[1] == self.embed_dim

        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)
        assert static_embedded.shape[1] == self.embed_dim

        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])
        assert x_unembed.shape[0] == batch_size
        assert x_unembed.shape[1] == self.global_shape_mu[0] * self.global_shape_mu[1]
        assert x_unembed.shape[2] == self.local_shape_mu[0] * self.local_shape_mu[1]
        assert (
            x_unembed.shape[3]
            == self.in_channels * self.patch_size_px[0] * self.patch_size_px[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


def masked_inference(
        model,
        batch: dict[str, torch.Tensor],
        mask_global: Optional[torch.Tensor] = None,
        mask_local: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    """
    Perform masked inference with Prithvi-WxC model

    Args:
        model: The PrithviWxC model to perform the inference with.
        batch: Dictionary containing the keys 'x', 'y', 'input_time',
            'lead_time' and 'static'. The associated torch tensors have the
            following shapes:
            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.
        mask_global: A 2D tensor of shape [n_lats_px / mask_unit_size_px[0], n_lons_px / patch_size_px[1]]
            identifying the global patches to mask.
        mask_local: A 2D tensor of shape [mask_unit_size_px[0] / patch_size_px[0], mask_unit_size_px[1] / patch_size_px[1]]
            identifying the global patches to mask.
    Returns:
        Tensor of shape [batch, parameter, lat, lon].
    """
    assert batch["x"].shape[2] == model.in_channels
    assert batch["x"].shape[3] == model.n_lats_px
    assert batch["x"].shape[4] == model.n_lons_px
    assert batch["y"].shape[1] == model.in_channels
    assert batch["y"].shape[2] == model.n_lats_px
    assert batch["y"].shape[3] == model.n_lons_px
    if model.positional_encoding == 'fourier':
        # the first two features (lat, lon) are encoded separately
        assert batch['static'].shape[1] - 2 == model.in_channels_static, "When setting model.positional_encoding to fourier, the number of static params change in the dataset. So, in the config, reduce num_static_channels (e.g., 4 instead of 7)."
    else:
        assert batch['static'].shape[1] == model.in_channels_static
    assert batch["static"].shape[2] == model.n_lats_px
    assert batch["static"].shape[3] == model.n_lons_px

    x_rescaled = (batch["x"] - model.input_scalers_mu) / (
        model.input_scalers_sigma + model.input_scalers_epsilon
    )
    batch_size = x_rescaled.shape[0]

    if model.positional_encoding == 'fourier':
        x_static_pos = model.fourier_pos_encoding(batch['static']) # B, embed_dim, lat / patch_size, lon / patch_size
        x_static = (batch['static'][:, 2:] - model.static_input_scalers_mu[:, 3:]) / ( # The first two channels in batch['static'] are used in positional encoding
            model.static_input_scalers_sigma[:, 3:] + model.static_input_scalers_epsilon # This translates to the first three channels in 'static_input_scalers_mu'
        )
    else:
        x_static = (batch["static"] - model.static_input_scalers_mu) / (
            model.static_input_scalers_sigma + model.static_input_scalers_epsilon
        )

    if model.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)
        assert (
            batch["y"].shape == x_hat.shape
        ), f'Shapes {batch["y"].shape} and {x_hat.shape} do not agree.'
    elif model.residual == "climate":
        climate_scaled = (
            batch["climate"] - model.input_scalers_mu.view(1, -1, 1, 1)
        ) / (
            model.input_scalers_sigma.view(1, -1, 1, 1) + model.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 = model.parameter_dropout(x_rescaled)

    x_embedded = model.patch_embedding(x_rescaled)
    assert x_embedded.shape[1] == model.embed_dim

    if model.residual == "climate":
        static_embedded = model.patch_embedding_static(
            torch.cat((x_static, climate_scaled), dim=1)
        )
    else:
        static_embedded = model.patch_embedding_static(x_static)
    assert static_embedded.shape[1] == model.embed_dim

    if model.positional_encoding == 'fourier':
        static_embedded += x_static_pos

    x_embedded = model.to_patching(x_embedded)
    static_embedded = model.to_patching(static_embedded)

    time_encoding = model.time_encoding(batch['input_time'], batch['lead_time'])

    tokens = x_embedded + static_embedded + time_encoding

    if mask_global is not None:
        indices_masked = torch.where(mask_global.flatten())[0]
        indices_masked = torch.repeat_interleave(indices_masked[None], tokens.shape[0], dim=0)
        indices_unmasked = torch.where(~mask_global.flatten())[0]
        indices_unmasked = torch.repeat_interleave(indices_unmasked[None], tokens.shape[0], dim=0)
    elif mask_local is not None:
        indices_masked = torch.where(mask_local.flatten())[0]
        indices_masked = indices_masked[None, None].expand(tokens.shape[0], tokens.shape[1], -1)
        indices_unmasked = torch.where(~mask_local.flatten())[0]
        indices_unmasked = indices_unmasked[None, None].expand(tokens.shape[0], tokens.shape[1], -1)
        print(indices_unmasked.shape)
    else:
        mask_global = torch.zeros(model.global_shape_mu, dtype=torch.bool)
        indices_masked = torch.where(mask_global.flatten())[0]
        indices_masked = torch.repeat_interleave(indices_masked[None], tokens.shape[0], dim=0)
        indices_unmasked = torch.where(~mask_global.flatten())[0]
        indices_unmasked = torch.repeat_interleave(indices_unmasked[None], tokens.shape[0], dim=0)

    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 = model.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 = model.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, _ = model.reconstruct_batch(
        indices_masked, indices_unmasked, masked, x_encoded
    )

    x_decoded = model.decoder(recon)

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

    # Reshape to (batch, global_lat, global_lon, local_lat, local_lon, in_channels * patch_size[0] * patch_size[1])
    assert x_unembed.shape[0] == batch_size
    assert x_unembed.shape[1] == model.global_shape_mu[0] * model.global_shape_mu[1]
    assert x_unembed.shape[2] == model.local_shape_mu[0] * model.local_shape_mu[1]
    assert (
        x_unembed.shape[3]
        == model.in_channels * model.patch_size_px[0] * model.patch_size_px[1]
    )

    x_out = model.from_patching(x_unembed)

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

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

    return x_out