model/aurora_network.py

"""Copyright (c) Microsoft Corporation. Licensed under the MIT license."""

import contextlib
import dataclasses
import warnings
from datetime import timedelta
from typing import Optional

import numpy as np
import torch
import torch.nn as nn
from huggingface_hub import hf_hub_download
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
    apply_activation_checkpointing,
)

from .aurora_batch import Batch
from .aurora_insolation import insolation
from .aurora_compat import (
    _adapt_checkpoint_air_pollution,
    _adapt_checkpoint_pretrained,
    _adapt_checkpoint_v1p5,
    _adapt_checkpoint_wave,
)
from .aurora_decoder import Perceiver3DDecoder
from .aurora_encoder import Perceiver3DEncoder
from .aurora_lora import LoRAMode
from .aurora_normalisation import log_transform, log_untransform
from .aurora_perceiver import PerceiverAttention
from .aurora_swin3d import Swin3DTransformerBackbone, WindowAttention

__all__ = [
    "Aurora",
    "AuroraPretrained",
    "AuroraSmallPretrained",
    "AuroraSmall",
    "Aurora12hPretrained",
    "AuroraHighRes",
    "AuroraAirPollution",
    "AuroraWave",
    "AuroraV1p5",
    "AuroraV1p5Ensemble",
]


class Aurora(torch.nn.Module):
    """The Aurora model.

    Defaults to the 1.3 B parameter configuration.

    Also supports ensemble forecasts.
    """

    default_checkpoint_repo = "microsoft/aurora"
    """str: Name of the HuggingFace repository to load the default checkpoint from."""

    default_checkpoint_name = "aurora-0.25-finetuned.ckpt"
    """str: Name of the default checkpoint."""

    default_checkpoint_revision = "0be7e57c685dac86b78c4a19a3ab149d13c6a3dd"
    """str: Commit hash of the default checkpoint."""

    def __init__(
        self,
        *,
        surf_vars: tuple[str, ...] = ("2t", "10u", "10v", "msl"),
        static_vars: tuple[str, ...] = ("lsm", "z", "slt"),
        atmos_vars: tuple[str, ...] = ("z", "u", "v", "t", "q"),
        window_size: tuple[int, int, int] = (2, 6, 12),
        encoder_depths: tuple[int, ...] = (6, 10, 8),
        encoder_num_heads: tuple[int, ...] = (8, 16, 32),
        decoder_depths: tuple[int, ...] = (8, 10, 6),
        decoder_num_heads: tuple[int, ...] = (32, 16, 8),
        latent_levels: int = 4,
        patch_size: int = 4,
        embed_dim: int = 512,
        num_heads: int = 16,
        mlp_ratio: float = 4.0,
        drop_path: float = 0.0,
        drop_rate: float = 0.0,
        enc_depth: int = 1,
        dec_depth: int = 1,
        dec_mlp_ratio: float = 2.0,
        perceiver_ln_eps: float = 1e-5,
        max_history_size: int = 2,
        timestep: timedelta = timedelta(hours=6),
        stabilise_level_agg: bool = False,
        use_lora: bool = True,
        lora_steps: int = 40,
        lora_mode: LoRAMode = "single",
        surf_stats: Optional[dict[str, tuple[float, float]]] = None,
        autocast: bool = False,
        autocast_dtype: torch.dtype = torch.bfloat16,
        bf16_mode: bool = False,
        use_fp16_safe_attention: bool = False,
        level_condition: Optional[tuple[int | float, ...]] = None,
        dynamic_vars: bool = False,
        atmos_static_vars: bool = False,
        separate_perceiver: tuple[str, ...] = (),
        modulation_heads: tuple[str, ...] = (),
        positive_surf_vars: tuple[str, ...] = (),
        positive_atmos_vars: tuple[str, ...] = (),
        clamp_at_first_step: bool = False,
        simulate_indexing_bug: bool = False,
        stochastic: bool = False,
        use_updated_lead_time_embedding: bool = False,
        variable_lead_time: bool = False,
        rollout_input_clipping: Optional[dict[str, dict[str, Optional[float]]]] = None,
        output_only_surf_vars: tuple[str, ...] = (),
        output_only_atmos_vars: tuple[str, ...] = (),
    ) -> None:
        """Construct an instance of the model.

        Args:
            surf_vars (tuple[str, ...], optional): All surface-level variables supported by the
                model.
            static_vars (tuple[str, ...], optional): All static variables supported by the
                model.
            atmos_vars (tuple[str, ...], optional): All atmospheric variables supported by the
                model.
            window_size (tuple[int, int, int], optional): Vertical height, height, and width of the
                window of the underlying Swin transformer.
            encoder_depths (tuple[int, ...], optional): Number of blocks in each encoder layer.
            encoder_num_heads (tuple[int, ...], optional): Number of attention heads in each encoder
                layer. The dimensionality doubles after every layer. To keep the dimensionality of
                every head constant, you want to double the number of heads after every layer. The
                dimensionality of attention head of the first layer is determined by `embed_dim`
                divided by the value here. For all cases except one, this is equal to `64`.
            decoder_depths (tuple[int, ...], optional): Number of blocks in each decoder layer.
                Generally, you want this to be the reversal of `encoder_depths`.
            decoder_num_heads (tuple[int, ...], optional): Number of attention heads in each decoder
                layer. Generally, you want this to be the reversal of `encoder_num_heads`.
            latent_levels (int, optional): Number of latent pressure levels.
            patch_size (int, optional): Patch size.
            embed_dim (int, optional): Patch embedding dimension.
            num_heads (int, optional): Number of attention heads in the aggregation and
                deaggregation blocks. The dimensionality of these attention heads will be equal to
                `embed_dim` divided by this value.
            mlp_ratio (float, optional): Hidden dim. to embedding dim. ratio for MLPs.
            drop_rate (float, optional): Drop-out rate.
            drop_path (float, optional): Drop-path rate.
            enc_depth (int, optional): Number of Perceiver blocks in the encoder.
            dec_depth (int, optional): Number of Perceiver blocks in the decoder.
            dec_mlp_ratio (float, optional): Hidden dim. to embedding dim. ratio for MLPs in the
                decoder. The embedding dimensionality here is different, which is why this is a
                separate parameter.
            perceiver_ln_eps (float, optional): Epsilon in the perceiver layer norm. layers. Used
                to stabilise the model.
            max_history_size (int, optional): Maximum number of history steps. You can load
                checkpoints with a smaller `max_history_size`, but you cannot load checkpoints
                with a larger `max_history_size`.
            timestep (timedelta, optional): Timestep of the model. Defaults to 6 hours.
            stabilise_level_agg (bool, optional): Stabilise the level aggregation by inserting an
                additional layer normalisation. Defaults to `False`.
            use_lora (bool, optional): Use LoRA adaptation.
            lora_steps (int, optional): Use different LoRA adaptation for the first so-many roll-out
                steps.
            lora_mode (str, optional): LoRA mode. `"single"` uses the same LoRA for all roll-out
                steps, `"from_second"` uses the same LoRA from the second roll-out step on, and
                `"all"` uses a different LoRA for every roll-out step. Defaults to `"single"`.
            surf_stats (dict[str, tuple[float, float]], optional): For these surface-level
                variables, adjust the normalisation to the given tuple consisting of a new location
                and scale.
            autocast (bool, optional): To reduce memory usage, `torch.autocast` only the backbone
                to a lower-precision dtype. This is critical to enable fine-tuning.
            autocast_dtype (torch.dtype, optional): Data type to use when `autocast` is enabled.
                Defaults to `torch.bfloat16`.
            use_fp16_safe_attention (bool, optional): Replace
                :func:`torch.nn.functional.scaled_dot_product_attention` with a manual
                implementation that clamps intermediate values to prevent float16 overflow.
                Recommended when running with `autocast_dtype=torch.float16`.
                Defaults to `False`.
            level_condition (tuple[int | float, ...], optional): Make the patch embeddings dependent
                on pressure level. If you want to enable this feature, provide a tuple of all
                possible pressure levels.
            dynamic_vars (bool, optional): Use dynamically generated static variables, like time
                of day. Defaults to `False`.
            atmos_static_vars (bool, optional): Also concatenate the static variables to the
                atmospheric variables. Defaults to `False`.
            separate_perceiver (tuple[str, ...], optional): In the decoder, use a separate Perceiver
                for specific atmospheric variables. This can be helpful at fine-tuning time to deal
                with variables that have a significantly different behaviour. If you want to enable
                this features, set this to the collection of variables that should be run on a
                separate Perceiver.
            modulation_heads (tuple[str, ...], optional): Names of every variable for which to
                enable an additional head, the so-called modulation head, that can be used to
                predict the difference.
            positive_surf_vars (tuple[str, ...], optional): Mark these surface-level variables as
                positive. Clamp them before running them through the encoder, and also clamp them
                when autoregressively rolling out the model. The variables are not clamped for the
                first roll-out step.
            positive_atmos_vars (tuple[str, ...], optional): Mark these atmospheric variables as
                positive. Clamp them before running them through the encoder, and also clamp them
                when autoregressively rolling out the model. The variables are not clamped for the
                first roll-out step.
            clamp_at_first_step (bool, optional): Clamp the positive variables for the first
                roll-out step. Should only be used for inference. Defaults to `False`.
            simulate_indexing_bug (bool, optional): Simulate an indexing bug that's present for the
                air pollution version of Aurora. This is necessary to obtain numerical equivalence
                to the original implementation. Defaults to `False`.
            stochastic (bool, optional): If `True`, enable stochastic mode with noise injection.
                Defaults to `False`.
            use_updated_lead_time_embedding (bool, optional): Whether to use the updated lead time
                embedding with a minimum wavelength of 2 hours. Defaults to `False`.
            variable_lead_time (bool, optional): If `True`, use per-sample lead times passed
                via the `lead_times` argument to `forward` (a tensor of shape `(batch,)` in
                hours) instead of the fixed `timestep`. When enabled, `lead_times` must be
                provided.
                Defaults to `False`.
            rollout_input_clipping (dict[str, dict[str, float]], optional): Per-variable
                clipping bounds applied to predictions during autoregressive rollout before they
                become the next input. Keys are variable names (must match `surf_vars` or
                `atmos_vars`). Values are dicts with optional `"min"` and `"max"` keys.
                Example: `{"tcc": {"min": 0, "max": 1}}`. Defaults to `None`.
            output_only_surf_vars (tuple[str, ...], optional): Surface-level variables that the
                model predicts but that are not present in real input data. These will be
                zero-padded in the input batch during rollout. Defaults to `()`.
            output_only_atmos_vars (tuple[str, ...], optional): Atmospheric variables that the
                model predicts but that are not present in real input data. These will be
                zero-padded in the input batch during rollout. Defaults to `()`.
        """
        super().__init__()
        self.surf_vars = surf_vars
        self.static_vars = static_vars
        self.atmos_vars = atmos_vars
        self.patch_size = patch_size
        self.surf_stats = surf_stats or dict()
        self.max_history_size = max_history_size
        self.timestep = timestep
        self.use_lora = use_lora
        self.positive_surf_vars = positive_surf_vars
        self.positive_atmos_vars = positive_atmos_vars
        self.clamp_at_first_step = clamp_at_first_step
        self.variable_lead_time = variable_lead_time
        self.rollout_input_clipping = rollout_input_clipping
        self.output_only_surf_vars = output_only_surf_vars
        self.output_only_atmos_vars = output_only_atmos_vars

        if self.surf_stats:
            warnings.warn(
                f"The normalisation statics for the following surface-level variables are manually "
                f"adjusted: {', '.join(sorted(self.surf_stats.keys()))}. "
                f"Please ensure that this is right!",
                stacklevel=2,
            )

        self.encoder = Perceiver3DEncoder(
            surf_vars=surf_vars,
            static_vars=static_vars,
            atmos_vars=atmos_vars,
            patch_size=patch_size,
            embed_dim=embed_dim,
            num_heads=num_heads,
            drop_rate=drop_rate,
            mlp_ratio=mlp_ratio,
            head_dim=embed_dim // num_heads,
            depth=enc_depth,
            latent_levels=latent_levels,
            max_history_size=max_history_size,
            perceiver_ln_eps=perceiver_ln_eps,
            stabilise_level_agg=stabilise_level_agg,
            level_condition=level_condition,
            dynamic_vars=dynamic_vars,
            atmos_static_vars=atmos_static_vars,
            simulate_indexing_bug=simulate_indexing_bug,
            use_updated_lead_time_embedding=use_updated_lead_time_embedding,
        )

        self.backbone = Swin3DTransformerBackbone(
            window_size=window_size,
            encoder_depths=encoder_depths,
            encoder_num_heads=encoder_num_heads,
            decoder_depths=decoder_depths,
            decoder_num_heads=decoder_num_heads,
            embed_dim=embed_dim,
            mlp_ratio=mlp_ratio,
            drop_path_rate=drop_path,
            attn_drop_rate=drop_rate,
            drop_rate=drop_rate,
            use_lora=use_lora,
            lora_steps=lora_steps,
            lora_mode=lora_mode,
            stochastic=stochastic,
            use_updated_lead_time_embedding=use_updated_lead_time_embedding,
        )

        self.decoder = Perceiver3DDecoder(
            surf_vars=surf_vars,
            atmos_vars=atmos_vars,
            patch_size=patch_size,
            # Concatenation at the backbone end doubles the dim.
            embed_dim=embed_dim * 2,
            head_dim=embed_dim * 2 // num_heads,
            num_heads=num_heads,
            depth=dec_depth,
            # Because of the concatenation, high ratios are expensive.
            # We use a lower ratio here to keep the memory in check.
            mlp_ratio=dec_mlp_ratio,
            perceiver_ln_eps=perceiver_ln_eps,
            level_condition=level_condition,
            separate_perceiver=separate_perceiver,
            modulation_heads=modulation_heads,
        )

        if bf16_mode and not autocast:
            warnings.warn(
                "`bf16_mode` was removed, because it caused serious issues for gradient "
                "computation. `bf16_mode` now automatically activates `autocast`, which will not "
                "save as much memory, but should be much more stable.",
                stacklevel=2,
            )
            autocast = True

        self.autocast = autocast
        self.autocast_dtype = autocast_dtype
        self.autocast_encoder = False
        self.autocast_backbone = autocast
        self.autocast_decoder = False

        # Enable fp16-safe attention on all attention modules.
        if use_fp16_safe_attention:
            for m in self.modules():
                if isinstance(m, (WindowAttention, PerceiverAttention)):
                    m.use_fp16_safe_attention = True

    def reset_noise(self) -> None:
        """Flush the backbone noise cache.

        See :meth:`Swin3DTransformerBackbone.reset_noise`."""
        self.backbone.reset_noise()

    def set_noise_accumulation(self, n: int = 0) -> None:
        """Enable or disable noise caching in the backbone.

        See :meth:`Swin3DTransformerBackbone.set_noise_accumulation`.

        Args:
            n (int): Number of steps for noise accumulation. Disables accumulation if `n=0`.
        """
        self.backbone.set_noise_accumulation(n)

    def forward(self, batch: Batch, lead_times: Optional[torch.Tensor] = None) -> Batch:
        """Forward pass.

        Args:
            batch (:class:`aurora.Batch`): Batch to run the model on.
            lead_times (:class:`torch.Tensor`, optional): Per-sample lead times of shape
                `(batch,)` in hours. Required when the model was configured with
                `variable_lead_time=True`. Ignored otherwise.

        Returns:
            :class:`Batch`: Prediction for the batch.
        """
        batch = self.batch_transform_hook(batch)

        # Get the first parameter. We'll derive the data type and device from this parameter.
        p = next(self.parameters())
        batch = batch.type(p.dtype)
        batch = self._pre_norm_hook(batch)
        batch = batch.normalise(surf_stats=self.surf_stats)
        batch = batch.crop(patch_size=self.patch_size)
        batch = batch.to(p.device)

        H, W = batch.spatial_shape
        patch_res = (
            self.encoder.latent_levels,
            H // self.encoder.patch_size,
            W // self.encoder.patch_size,
        )

        # Insert batch and history dimension for static variables.
        B, T = next(iter(batch.surf_vars.values())).shape[:2]
        batch = dataclasses.replace(
            batch,
            static_vars={k: v[None, None].repeat(B, T, 1, 1) for k, v in batch.static_vars.items()},
        )

        # Apply some transformations before feeding `batch` to the encoder. We'll later want to
        # refer to the original batch too, so rename the variable.
        transformed_batch = batch

        # Clamp positive variables.
        if self.positive_surf_vars:
            transformed_batch = dataclasses.replace(
                transformed_batch,
                surf_vars={
                    k: v.clamp(min=0) if k in self.positive_surf_vars else v
                    for k, v in batch.surf_vars.items()
                },
            )
        if self.positive_atmos_vars:
            transformed_batch = dataclasses.replace(
                transformed_batch,
                atmos_vars={
                    k: v.clamp(min=0) if k in self.positive_atmos_vars else v
                    for k, v in batch.atmos_vars.items()
                },
            )

        transformed_batch = self._pre_encoder_hook(transformed_batch)

        # Resolve lead times to a tensor of shape (B,) in hours.
        if self.variable_lead_time:
            if lead_times is None:
                raise ValueError(
                    "`variable_lead_time=True` but `lead_times` is `None`. "
                    "Please provide a `lead_times` tensor of shape `(batch,)` in hours."
                )
            lead_times = lead_times.to(device=p.device, dtype=p.dtype)
        else:
            lead_hours = self.timestep.total_seconds() / 3600
            lead_times = torch.full((B,), lead_hours, device=p.device, dtype=p.dtype)

        if torch.cuda.is_available():
            device_type = "cuda"
        elif torch.xpu.is_available():
            device_type = "xpu"
        else:
            device_type = "cpu"
        autocast = torch.autocast(device_type=device_type, dtype=self.autocast_dtype)
        context_encoder = autocast if self.autocast_encoder else contextlib.nullcontext()
        context_backbone = autocast if self.autocast_backbone else contextlib.nullcontext()
        context_decoder = autocast if self.autocast_decoder else contextlib.nullcontext()

        with context_encoder:
            x = self.encoder(
                transformed_batch,
                lead_times=lead_times,
            )
        with context_backbone:
            x = self.backbone(
                x,
                lead_times=lead_times,
                patch_res=patch_res,
                rollout_step=batch.metadata.rollout_step,
            )
        with context_decoder:
            pred = self.decoder(
                x,
                batch,
                lead_times=lead_times,
                patch_res=patch_res,
            )

        # Remove batch and history dimension from static variables.
        pred = dataclasses.replace(
            pred,
            static_vars={k: v[0, 0] for k, v in batch.static_vars.items()},
        )

        # Insert history dimension in prediction. The time should already be right.
        pred = dataclasses.replace(
            pred,
            surf_vars={k: v[:, None] for k, v in pred.surf_vars.items()},
            atmos_vars={k: v[:, None] for k, v in pred.atmos_vars.items()},
        )

        pred = self._post_decoder_hook(batch, pred)

        # Clamp positive variables.
        clamp_at_rollout_step = (
            pred.metadata.rollout_step >= 1
            if self.clamp_at_first_step
            else pred.metadata.rollout_step > 1
        )
        if self.positive_surf_vars and clamp_at_rollout_step:
            pred = dataclasses.replace(
                pred,
                surf_vars={
                    k: v.clamp(min=0) if k in self.positive_surf_vars else v
                    for k, v in pred.surf_vars.items()
                },
            )
        if self.positive_atmos_vars and clamp_at_rollout_step:
            pred = dataclasses.replace(
                pred,
                atmos_vars={
                    k: v.clamp(min=0) if k in self.positive_atmos_vars else v
                    for k, v in pred.atmos_vars.items()
                },
            )

        # Cast to float32 before unnormalising to avoid overflow.
        pred = pred.type(torch.float32)
        pred = pred.unnormalise(surf_stats=self.surf_stats)

        pred = self._post_unnorm_hook(batch, pred)

        return pred

    def batch_transform_hook(self, batch: Batch) -> Batch:
        """Transform the batch right after receiving it and before normalisation.

        This function should be idempotent.
        """
        return batch

    def _pre_encoder_hook(self, batch: Batch) -> Batch:
        """Transform the batch before it goes through the encoder."""
        return batch

    def _pre_norm_hook(self, batch: Batch) -> Batch:
        """Transform the batch before normalisation.

        This is called automatically in :meth:`forward` right before :meth:`Batch.normalise`. Unlike
        :meth:`batch_transform_hook`, this hook is *not* called separately in rollout, so
        non-idempotent transforms (e.g. log-scaling) belong here.
        """
        return batch

    def _post_decoder_hook(self, batch: Batch, pred: Batch) -> Batch:
        """Transform the prediction right after the decoder."""
        return pred

    def _post_unnorm_hook(self, batch: Batch, pred: Batch) -> Batch:
        """Transform the prediction after un-normalisation, in physical space.

        Subclasses can override this to apply post-processing that must operate on un-normalised
        (physical) values, such as inverse log-scaling or recomputing prescribed channels.
        """
        return pred

    def apply_rollout_input_clipping(self, pred: Batch) -> Batch:
        """Clamp specified variables according to `rollout_input_clipping`.

        This is intended to be called during autoregressive rollout *before* feeding a prediction
        back as input, so that the unclipped prediction remains available for loss computation
        during training. To minimize any other changes to the data flow from models prior to V1p5,
        this is not called automatically in :meth:`forward`.
        """
        if not self.rollout_input_clipping:
            return pred

        clipped_surf = dict(pred.surf_vars)
        clipped_atmos = dict(pred.atmos_vars)

        for var_name, bounds in self.rollout_input_clipping.items():
            lo = bounds.get("min")
            hi = bounds.get("max")
            if var_name in clipped_surf:
                v = clipped_surf[var_name]
                if lo is not None:
                    v = v.clamp(min=lo)
                if hi is not None:
                    v = v.clamp(max=hi)
                clipped_surf[var_name] = v
            if var_name in clipped_atmos:
                v = clipped_atmos[var_name]
                if lo is not None:
                    v = v.clamp(min=lo)
                if hi is not None:
                    v = v.clamp(max=hi)
                clipped_atmos[var_name] = v

        return dataclasses.replace(pred, surf_vars=clipped_surf, atmos_vars=clipped_atmos)

    def load_checkpoint(
        self,
        repo: Optional[str] = None,
        name: Optional[str] = None,
        revision: Optional[str] = None,
        strict: bool = True,
    ) -> None:
        """Load a checkpoint from HuggingFace.

        Args:
            repo (str, optional): Name of the repository of the form `user/repo`.
            name (str, optional): Path to the checkpoint relative to the root of the repository,
                e.g. `checkpoint.cpkt`.
            revision (str, optional): Version hash of the Huggingface git repository commit.
            strict (bool, optional): Error if the model parameters are not exactly equal to the
                parameters in the checkpoint. Defaults to `True`.
        """
        repo = repo or self.default_checkpoint_repo
        name = name or self.default_checkpoint_name
        revision = revision or self.default_checkpoint_revision
        path = hf_hub_download(repo_id=repo, filename=name, revision=revision)
        self.load_checkpoint_local(path, strict=strict)

    def load_checkpoint_local(self, path: str, strict: bool = True) -> None:
        """Load a checkpoint directly from a file.

        Args:
            path (str): Path to the checkpoint.
            strict (bool, optional): Error if the model parameters are not exactly equal to the
                parameters in the checkpoint. Defaults to `True`.
        """
        # Assume that all parameters are either on the CPU or on the GPU.
        device = next(self.parameters()).device
        d = torch.load(path, map_location=device, weights_only=True)

        d = self._adapt_checkpoint(d)

        # Check if the history size is compatible and adjust weights if necessary.
        current_history_size = d["encoder.surf_token_embeds.weights.2t"].shape[2]
        if self.max_history_size > current_history_size:
            self.adapt_checkpoint_max_history_size(d)
        elif self.max_history_size < current_history_size:
            raise AssertionError(
                f"Cannot load checkpoint with `max_history_size` {current_history_size} "
                f"into model with `max_history_size` {self.max_history_size}."
            )

        self.load_state_dict(d, strict=strict)

    def _adapt_checkpoint(self, d: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
        """Adapt an existing checkpoint to make it compatible with the current version of the model.

        Args:
            d (dict[str, torch.Tensor]): Checkpoint.

        Return:
            dict[str, torch.Tensor]: Adapted checkpoint.
        """
        return _adapt_checkpoint_pretrained(self.patch_size, d)

    def adapt_checkpoint_max_history_size(self, checkpoint: dict[str, torch.Tensor]) -> None:
        """Adapt a checkpoint with smaller `max_history_size` to a model with a larger
        `max_history_size` than the current model.

        If a checkpoint was trained with a larger `max_history_size` than the current model,
        this function will assert fail to prevent loading the checkpoint. This is to
        prevent loading a checkpoint which will likely cause the checkpoint to degrade its
        performance.

        This implementation copies weights from the checkpoint to the model and fills zeros
        for the new history width dimension. It mutates `checkpoint`.
        """
        for name, weight in list(checkpoint.items()):
            # We only need to adapt the patch embedding in the encoder.
            enc_surf_embedding = name.startswith("encoder.surf_token_embeds.weights.")
            enc_atmos_embedding = name.startswith("encoder.atmos_token_embeds.weights.")
            if enc_surf_embedding or enc_atmos_embedding:
                # This shouldn't get called with current logic but leaving here for future proofing
                # and in cases where its called outside current context.
                if not (weight.shape[2] <= self.max_history_size):
                    raise AssertionError(
                        f"Cannot load checkpoint with `max_history_size` {weight.shape[2]} "
                        f"into model with `max_history_size` {self.max_history_size}."
                    )

                # Initialize the new weight tensor.
                new_weight = torch.zeros(
                    (weight.shape[0], 1, self.max_history_size, weight.shape[3], weight.shape[4]),
                    device=weight.device,
                    dtype=weight.dtype,
                )
                # Copy the existing weights to the new tensor by duplicating the histories provided
                # into any new history dimensions. The rest remains at zero.
                new_weight[:, :, : weight.shape[2]] = weight

                checkpoint[name] = new_weight

    def configure_activation_checkpointing(
        self,
        module_names: tuple[str, ...] = (
            "Basic3DDecoderLayer",
            "Basic3DEncoderLayer",
            "LinearPatchReconstruction",
            "Perceiver3DDecoder",
            "Perceiver3DEncoder",
            "Swin3DTransformerBackbone",
            "Swin3DTransformerBlock",
        ),
    ) -> None:
        """Configure activation checkpointing.

        This is required in order to compute gradients without running out of memory.

        Args:
            module_names (tuple[str, ...], optional): Names of the modules to checkpoint
                on.

        Raises:
            RuntimeError: If any module specifies in `module_names` was not found and
                thus could not be checkpointed.
        """

        found: set[str] = set()

        def check(x: torch.nn.Module) -> bool:
            name = x.__class__.__name__
            if name in module_names:
                found.add(name)
                return True
            else:
                return False

        apply_activation_checkpointing(self, check_fn=check)

        if found != set(module_names):
            raise RuntimeError(
                f"Could not checkpoint on the following modules: "
                f"{', '.join(sorted(set(module_names) - found))}."
            )


class AuroraPretrained(Aurora):
    """Pretrained version of Aurora."""

    default_checkpoint_name = "aurora-0.25-pretrained.ckpt"
    default_checkpoint_revision = "0be7e57c685dac86b78c4a19a3ab149d13c6a3dd"

    def __init__(
        self,
        *,
        use_lora: bool = False,
        **kw_args,
    ) -> None:
        super().__init__(
            use_lora=use_lora,
            **kw_args,
        )


class AuroraSmallPretrained(Aurora):
    """Small pretrained version of Aurora.

    Should only be used for debugging.
    """

    default_checkpoint_name = "aurora-0.25-small-pretrained.ckpt"
    default_checkpoint_revision = "0be7e57c685dac86b78c4a19a3ab149d13c6a3dd"

    def __init__(
        self,
        *,
        encoder_depths: tuple[int, ...] = (2, 6, 2),
        encoder_num_heads: tuple[int, ...] = (4, 8, 16),
        decoder_depths: tuple[int, ...] = (2, 6, 2),
        decoder_num_heads: tuple[int, ...] = (16, 8, 4),
        embed_dim: int = 256,
        num_heads: int = 8,
        use_lora: bool = False,
        **kw_args,
    ) -> None:
        super().__init__(
            encoder_depths=encoder_depths,
            encoder_num_heads=encoder_num_heads,
            decoder_depths=decoder_depths,
            decoder_num_heads=decoder_num_heads,
            embed_dim=embed_dim,
            num_heads=num_heads,
            use_lora=use_lora,
            **kw_args,
        )


AuroraSmall = AuroraSmallPretrained  #: Alias for backwards compatibility


class Aurora12hPretrained(Aurora):
    """Pretrained version of Aurora with time step 12 hours."""

    default_checkpoint_name = "aurora-0.25-12h-pretrained.ckpt"
    default_checkpoint_revision = "15e76e47b65bf4b28fd2246b7b5b951d6e2443b9"

    def __init__(
        self,
        *,
        timestep: timedelta = timedelta(hours=12),
        use_lora: bool = False,
        **kw_args,
    ) -> None:
        super().__init__(
            timestep=timestep,
            use_lora=use_lora,
            **kw_args,
        )


class AuroraHighRes(Aurora):
    """High-resolution version of Aurora."""

    default_checkpoint_name = "aurora-0.1-finetuned.ckpt"
    default_checkpoint_revision = "0be7e57c685dac86b78c4a19a3ab149d13c6a3dd"

    def __init__(
        self,
        *,
        patch_size: int = 10,
        encoder_depths: tuple[int, ...] = (6, 8, 8),
        decoder_depths: tuple[int, ...] = (8, 8, 6),
        **kw_args,
    ) -> None:
        super().__init__(
            patch_size=patch_size,
            encoder_depths=encoder_depths,
            decoder_depths=decoder_depths,
            **kw_args,
        )


class AuroraAirPollution(Aurora):
    """Fine-tuned version of Aurora for air pollution."""

    default_checkpoint_name = "aurora-0.4-air-pollution.ckpt"
    default_checkpoint_revision = "1764d5630a53d3d7a7d169ca335236fc343e4bfc"

    _predict_difference_history_dim_lookup = {
        "pm1": 0,
        "pm2p5": 0,
        "pm10": 0,
        "co": 1,
        "tcco": 1,
        "no": 0,
        "tc_no": 0,
        "no2": 0,
        "tcno2": 0,
        "so2": 1,
        "tcso2": 1,
        "go3": 1,
        "gtco3": 1,
    }
    """dict[str, int]: For every variable that we want to predict the difference for, the index
    into the history dimension that should be used when predicting the difference."""

    def __init__(
        self,
        *,
        surf_vars: tuple[str, ...] = (
            ("2t", "10u", "10v", "msl")
            + ("pm1", "pm2p5", "pm10", "tcco", "tc_no", "tcno2", "gtco3", "tcso2")
        ),
        static_vars: tuple[str, ...] = (
            ("lsm", "z", "slt")
            + ("static_ammonia", "static_ammonia_log", "static_co", "static_co_log")
            + ("static_nox", "static_nox_log", "static_so2", "static_so2_log")
        ),
        atmos_vars: tuple[str, ...] = ("z", "u", "v", "t", "q", "co", "no", "no2", "go3", "so2"),
        patch_size: int = 3,
        timestep: timedelta = timedelta(hours=12),
        level_condition: Optional[tuple[int | float, ...]] = (
            (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)
        ),
        dynamic_vars: bool = True,
        atmos_static_vars: bool = True,
        separate_perceiver: tuple[str, ...] = ("co", "no", "no2", "go3", "so2"),
        modulation_heads: tuple[str, ...] = tuple(_predict_difference_history_dim_lookup.keys()),
        positive_surf_vars: tuple[str, ...] = (
            ("pm1", "pm2p5", "pm10", "tcco", "tc_no", "tcno2", "gtco3", "tcso2")
        ),
        positive_atmos_vars: tuple[str, ...] = ("co", "no", "no2", "go3", "so2"),
        simulate_indexing_bug: bool = True,
        **kw_args,
    ) -> None:
        super().__init__(
            surf_vars=surf_vars,
            static_vars=static_vars,
            atmos_vars=atmos_vars,
            patch_size=patch_size,
            timestep=timestep,
            level_condition=level_condition,
            dynamic_vars=dynamic_vars,
            atmos_static_vars=atmos_static_vars,
            separate_perceiver=separate_perceiver,
            modulation_heads=modulation_heads,
            positive_surf_vars=positive_surf_vars,
            positive_atmos_vars=positive_atmos_vars,
            simulate_indexing_bug=simulate_indexing_bug,
            **kw_args,
        )

        self.surf_feature_combiner = torch.nn.ParameterDict(
            {v: nn.Linear(2, 1, bias=True) for v in self.positive_surf_vars}
        )
        self.atmos_feature_combiner = torch.nn.ParameterDict(
            {v: nn.Linear(2, 1, bias=True) for v in self.positive_atmos_vars}
        )
        for p in (*self.surf_feature_combiner.values(), *self.atmos_feature_combiner.values()):
            nn.init.constant_(p.weight, 0.5)
            nn.init.zeros_(p.bias)

    def _pre_encoder_hook(self, batch: Batch) -> Batch:
        # Transform the spikey variables with a specific log-transform before feeding them
        # to the encoder. See the paper for a motivation for the precise form of the transform.

        eps = 1e-4
        divisor = -np.log(eps)

        def _transform(z: torch.Tensor, feature_combiner: nn.Module) -> torch.Tensor:
            return feature_combiner(
                torch.stack(
                    [
                        z.clamp(min=0, max=2.5),
                        (torch.log(z.clamp(min=eps)) - np.log(eps)) / divisor,
                    ],
                    dim=-1,
                )
            )[..., 0]

        return dataclasses.replace(
            batch,
            surf_vars={
                k: _transform(v, self.surf_feature_combiner[k])
                if k in self.surf_feature_combiner
                else v
                for k, v in batch.surf_vars.items()
            },
            atmos_vars={
                k: _transform(v, self.atmos_feature_combiner[k])
                if k in self.atmos_feature_combiner
                else v
                for k, v in batch.atmos_vars.items()
            },
        )

    def _post_decoder_hook(self, batch: Batch, pred: Batch) -> Batch:
        # For this version of the model, we predict the difference. Specifically w.r.t. which
        # previous timestep (12 hours ago or 24 hours ago) is given by
        # `Aurora._predict_difference_history_dim_lookup`.

        dim_lookup = AuroraAirPollution._predict_difference_history_dim_lookup

        def _transform(
            prev: dict[str, torch.Tensor],
            model: dict[str, torch.Tensor],
            name: str,
        ) -> torch.Tensor:
            if name in dim_lookup:
                return model[name] + (1 + model[f"{name}_mod"]) * prev[name][:, dim_lookup[name]]
            else:
                return model[name]

        pred = dataclasses.replace(
            pred,
            surf_vars={k: _transform(batch.surf_vars, pred.surf_vars, k) for k in batch.surf_vars},
            atmos_vars={
                k: _transform(batch.atmos_vars, pred.atmos_vars, k) for k in batch.atmos_vars
            },
        )

        # When using LoRA, the lower-atmospheric levels of SO2 can be problematic and blow up.
        # We attempt to fix that by some very aggressive output clipping.
        if self.use_lora:
            parts: list[torch.Tensor] = []
            for i, level in enumerate(pred.metadata.atmos_levels):
                section = pred.atmos_vars["so2"][..., i, :, :]
                if level >= 850:
                    section = section.clamp(max=1)
                parts.append(section)
            pred.atmos_vars["so2"] = torch.stack(parts, dim=-3)

        return pred

    def _adapt_checkpoint(self, d: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
        d = Aurora._adapt_checkpoint(self, d)
        d = _adapt_checkpoint_air_pollution(self.patch_size, d)
        return d


class AuroraWave(Aurora):
    """Version of Aurora fined-tuned to HRES-WAM ocean wave data."""

    default_checkpoint_name = "aurora-0.25-wave.ckpt"
    default_checkpoint_revision = "74598e8c65d53a96077c08bb91acdfa5525340c9"

    def __init__(
        self,
        *,
        surf_vars: tuple[str, ...] = (
            ("2t", "10u", "10v", "msl")
            + ("swh", "mwd", "mwp", "pp1d", "shww", "mdww", "mpww", "shts", "mdts", "mpts")
            + ("swh1", "mwd1", "mwp1", "swh2", "mwd2", "mwp2", "wind", "10u_wave", "10v_wave")
        ),
        static_vars: tuple[str, ...] = ("lsm", "z", "slt", "wmb", "lat_mask"),
        lora_mode: LoRAMode = "from_second",
        stabilise_level_agg: bool = True,
        density_channel_surf_vars: tuple[str, ...] = (
            ("swh", "mwd", "mwp", "pp1d", "shww", "mdww", "mpww", "shts", "mdts", "mpts")
            + ("swh1", "mwd1", "mwp1", "swh2", "mwd2", "mwp2", "wind", "10u_wave", "10v_wave")
        ),
        angle_surf_vars: tuple[str, ...] = ("mwd", "mdww", "mdts", "mwd1", "mwd2"),
        **kw_args,
    ) -> None:
        # Model the density, sine, and cosine versions of the variables.
        supplemented_surf_vars: tuple[str, ...] = ()
        for name in surf_vars:
            if name in angle_surf_vars:
                supplemented_surf_vars += (f"{name}_sin", f"{name}_cos")
            else:
                supplemented_surf_vars += (name,)
            if name in density_channel_surf_vars:
                supplemented_surf_vars += (f"{name}_density",)

        super().__init__(
            surf_vars=supplemented_surf_vars,
            static_vars=static_vars,
            lora_mode=lora_mode,
            stabilise_level_agg=stabilise_level_agg,
            **kw_args,
        )

        self.density_channel_surf_vars = density_channel_surf_vars
        self.angle_surf_vars = angle_surf_vars

    def _adapt_checkpoint(self, d: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
        d = Aurora._adapt_checkpoint(self, d)
        d = _adapt_checkpoint_wave(self.patch_size, d)
        return d

    def batch_transform_hook(self, batch: Batch) -> Batch:
        #  Below we mutate `batch`, so make a copy here.
        batch = dataclasses.replace(batch, surf_vars=dict(batch.surf_vars))

        # It is important that these components are split off _before_ normalisation, as they
        # have specific normalisation statistics.
        if "dwi" in batch.surf_vars and "wind" in batch.surf_vars:
            # Split into u-component and v-component.
            u_wave = -batch.surf_vars["wind"] * torch.sin(torch.deg2rad(batch.surf_vars["dwi"]))
            v_wave = -batch.surf_vars["wind"] * torch.cos(torch.deg2rad(batch.surf_vars["dwi"]))

            # Update batch and remove `dwi`.
            batch.surf_vars["10u_wave"] = u_wave
            batch.surf_vars["10v_wave"] = v_wave
            del batch.surf_vars["dwi"]

        # If the magnitude of a wave is zero (or practically zero), it is absent, so indicate that
        # with NaNs. Only do this when data is given to the model and not when it is rolled out.
        if batch.metadata.rollout_step == 0:
            for name_sh, other_wave_components in [
                ("swh", ("mwd", "mwp", "pp1d")),
                ("shww", ("mdww", "mpww")),
                ("shts", ("mdts", "mdts")),
                ("swh1", ("mwd1", "mwp1")),
                ("swh2", ("mwd2", "mwp2")),
            ]:
                mask = batch.surf_vars[name_sh] < 1e-4
                if mask.sum() > 0:
                    for name in (name_sh,) + other_wave_components:
                        x = batch.surf_vars[name].clone()  # Clone to safely mutate.
                        x[mask] = np.nan
                        batch.surf_vars[name] = x
                        # There should be no small values left, except for in wave directions.
                        if name not in {"mwd", "mdww", "mdts", "mwd1", "mwd2"}:
                            assert (batch.surf_vars[name] < 1e-4).sum() == 0

        return batch

    def _pre_encoder_hook(self, batch: Batch) -> Batch:
        for name in list(batch.surf_vars):
            x = batch.surf_vars[name]

            # Create a density channel.
            if name in self.density_channel_surf_vars and f"{name}_density" not in batch.surf_vars:
                batch.surf_vars[f"{name}_density"] = (~torch.isnan(x)).float()
                batch.surf_vars[name] = x.nan_to_num(0)

            # Add sine and cosine values of the angle and remove the original angle variable
            sin_cos_present = f"{name}_sin" in batch.surf_vars and f"{name}_cos" in batch.surf_vars
            if name in self.angle_surf_vars and not sin_cos_present:
                batch.surf_vars[f"{name}_sin"] = torch.sin(torch.deg2rad(x)).nan_to_num(0)
                batch.surf_vars[f"{name}_cos"] = torch.cos(torch.deg2rad(x)).nan_to_num(0)
                del batch.surf_vars[name]

        return batch

    def _post_decoder_hook(self, batch: Batch, pred: Batch) -> Batch:
        wmb_mask = pred.static_vars["wmb"] > 0

        # Undo the sine and cosine components.
        for name in self.angle_surf_vars:
            if f"{name}_sin" in pred.surf_vars and f"{name}_cos" in pred.surf_vars:
                sin = pred.surf_vars[f"{name}_sin"]
                cos = pred.surf_vars[f"{name}_cos"]
                pred.surf_vars[name] = torch.rad2deg(torch.atan2(sin, cos)) % 360
                del pred.surf_vars[f"{name}_sin"]
                del pred.surf_vars[f"{name}_cos"]

        # Undo the density channels. First transform by a sigmoid to get the actual value of the
        # density channel.
        for name in self.density_channel_surf_vars:
            if name in pred.surf_vars:
                density = torch.sigmoid(pred.surf_vars[f"{name}_density"]) * wmb_mask
                data = pred.surf_vars[name] * wmb_mask
                data[density < 0.5] = np.nan
                pred.surf_vars[name] = data
                del pred.surf_vars[f"{name}_density"]

        return pred


class AuroraV1p5(Aurora):
    """Aurora 1.5 with expanded surface variables, variable lead-time support, and insolation.

    This variant was trained with an extended set of surface variables (26 total), additional static
    fields, and prescribed solar insolation as an input channel. It supports variable lead-time
    embeddings, enabling sub-6-hour prediction steps. Seven surface variables are output-only (not
    present in the real input data) and are zero-padded during autoregressive rollout.
    """

    default_checkpoint_repo = "ikwessel/aurora-1.5"
    default_checkpoint_name = "aurora-0.25-v1.5.ckpt"
    default_checkpoint_revision = "9751bb56e8e4a0f0a780e3cbe978f4c721e12bc7"

    def __init__(
        self,
        *,
        surf_vars: tuple[str, ...] = (
            ("2t", "10u", "10v", "msl", "2d", "tcwv", "tcc", "100u", "100v", "sp", "lcc", "mcc")
            + ("hcc", "skt", "stl1", "swvl1", "ci", "scaled_sd", "i10fg", "blh", "uvb_1h")
            + ("ssrd_1h", "ttr_1h", "scaled_tp_1h", "scaled_sf_1h", "insolation")
        ),
        static_vars: tuple[str, ...] = (
            ("lsm", "z", "anor", "isor", "cvh", "cl", "dl", "cvl", "slor", "slt_0", "slt_1")
            + ("slt_2", "slt_3", "slt_4", "slt_5", "slt_6", "slt_7", "sdfor", "sdor", "tvh_0")
            + ("tvh_18", "tvh_19", "tvh_3", "tvh_4", "tvh_5", "tvh_6", "tvl_0", "tvl_1", "tvl_10")
            + ("tvl_11", "tvl_13", "tvl_16", "tvl_17", "tvl_2", "tvl_7", "tvl_9")
        ),
        atmos_vars: tuple[str, ...] = ("z", "u", "v", "t", "q"),
        output_only_surf_vars: tuple[str, ...] = (
            ("i10fg", "blh", "uvb_1h", "ssrd_1h", "ttr_1h", "scaled_tp_1h", "scaled_sf_1h")
        ),
        rollout_input_clipping: Optional[dict[str, dict[str, Optional[float]]]] = None,
        variable_lead_time: bool = True,
        use_updated_lead_time_embedding: bool = True,
        use_lora: bool = False,
        use_fp16_safe_attention: bool = True,
        autocast: bool = True,
        autocast_dtype: torch.dtype = torch.float16,
        **kw_args,
    ) -> None:
        # Define default clipping ranges for rollout inputs, which can be overridden by passing in
        # `rollout_input_clipping`.
        rollout_input_clipping = rollout_input_clipping or {}
        if "tcwv" not in rollout_input_clipping:
            rollout_input_clipping["tcwv"] = {"min": 0.0, "max": None}
        if "tcc" not in rollout_input_clipping:
            rollout_input_clipping["tcc"] = {"min": 0.0, "max": 1.0}
        if "lcc" not in rollout_input_clipping:
            rollout_input_clipping["lcc"] = {"min": 0.0, "max": 1.0}
        if "mcc" not in rollout_input_clipping:
            rollout_input_clipping["mcc"] = {"min": 0.0, "max": 1.0}
        if "hcc" not in rollout_input_clipping:
            rollout_input_clipping["hcc"] = {"min": 0.0, "max": 1.0}
        if "swvl1" not in rollout_input_clipping:
            rollout_input_clipping["swvl1"] = {"min": 0.0, "max": 70.0}
        if "ci" not in rollout_input_clipping:
            rollout_input_clipping["ci"] = {"min": 0.0, "max": 1.0}
        if "scaled_sd" not in rollout_input_clipping:
            rollout_input_clipping["scaled_sd"] = {"min": 0.0, "max": 10.0}

        super().__init__(
            surf_vars=surf_vars,
            static_vars=static_vars,
            atmos_vars=atmos_vars,
            output_only_surf_vars=output_only_surf_vars,
            rollout_input_clipping=rollout_input_clipping,
            variable_lead_time=variable_lead_time,
            use_updated_lead_time_embedding=use_updated_lead_time_embedding,
            use_lora=use_lora,
            use_fp16_safe_attention=use_fp16_safe_attention,
            autocast=autocast,
            autocast_dtype=autocast_dtype,
            **kw_args,
        )
        self.autocast_encoder = autocast
        self.autocast_backbone = autocast
        self.autocast_decoder = autocast

        # Variable naming scheme assumes that all log-transformed variables start with "scaled_".
        self.log_transformed_surf_vars = tuple(v for v in self.surf_vars if v.startswith("scaled_"))

    def _pre_encoder_hook(self, batch: Batch) -> Batch:
        """Zero-pad output-only variables.

        Output-only variables are predicted by the model but are not present in real input data.
        They are added as zero tensors so the encoder receives the correct number of channels.
        Mutates `batch.surf_vars` / `batch.atmos_vars` in place so that both `batch` and
        `transformed_batch` in the caller see the new keys. Zero tensors are added post-
        normalization.
        """
        for var in self.output_only_surf_vars:
            ref = next(iter(batch.surf_vars.values()))
            batch.surf_vars[var] = torch.zeros_like(ref)
        for var in self.output_only_atmos_vars:
            ref = next(iter(batch.atmos_vars.values()))
            batch.atmos_vars[var] = torch.zeros_like(ref)
        return batch

    def _pre_norm_hook(self, batch: Batch) -> Batch:
        """Apply log-transform to scaled surface variables before normalisation."""
        return dataclasses.replace(
            batch,
            surf_vars={
                k: log_transform(v) if k in self.log_transformed_surf_vars else v
                for k, v in batch.surf_vars.items()
            },
        )

    def _post_unnorm_hook(self, batch: Batch, pred: Batch) -> Batch:
        """Apply inverse log-transform and recompute prescribed insolation."""
        pred = dataclasses.replace(
            pred,
            surf_vars={
                k: log_untransform(v) if k in self.log_transformed_surf_vars else v
                for k, v in pred.surf_vars.items()
            },
        )
        pred = self._update_insolation(pred)
        return pred

    def _update_insolation(self, pred: Batch) -> Batch:
        """Recompute prescribed insolation for the prediction's valid time."""
        if "insolation" not in pred.surf_vars:
            return pred

        lat_np = pred.metadata.lat.cpu().numpy().astype(np.float32)
        lon_np = pred.metadata.lon.cpu().numpy().astype(np.float32)

        sol_all = []
        for t in pred.metadata.time:
            sol = insolation([t], lat_np, lon_np, enforce_2d=True)
            sol_all.append(sol[0])  # Shape (H, W)
        sol_tensor = torch.tensor(
            np.stack(sol_all, axis=0),
            dtype=pred.surf_vars["insolation"].dtype,
            device=pred.surf_vars["insolation"].device,
        )
        # `pred.surf_vars["insolation"]` has shape (B, 1, H, W)
        sol_tensor = sol_tensor[:, None, :, :]

        return dataclasses.replace(
            pred,
            surf_vars={
                k: (sol_tensor if k == "insolation" else v) for k, v in pred.surf_vars.items()
            },
        )

    def _adapt_checkpoint(self, d: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
        return _adapt_checkpoint_v1p5(
            self.patch_size,
            self.surf_vars,
            self.static_vars,
            self.atmos_vars,
            d,
        )


class AuroraV1p5Ensemble(AuroraV1p5):
    """Aurora 1.5 ensemble version with stochastic noise injection."""

    default_checkpoint_name = "aurora-0.25-v1.5-ensemble.ckpt"
    default_checkpoint_revision = "9751bb56e8e4a0f0a780e3cbe978f4c721e12bc7"

    def __init__(
        self,
        *,
        stochastic: bool = True,
        **kw_args,
    ) -> None:
        super().__init__(
            stochastic=stochastic,
            **kw_args,
        )