Schemas

class anemoi.models.schemas.models.DefinedModels(value)

Bases: str, Enum

class anemoi.models.schemas.models.Model(*, _target_: DefinedModels, hidden_nodes_name: str | list[str], latent_skip: bool = True, _convert_: str = 'none')

Bases: BaseModel

target_: DefinedModels

Model object defined in anemoi.models.model.

hidden_nodes_name: str | list[str]

Name of the hidden nodes. If the model is hierarchical, it can be a list of names for each level.

latent_skip: bool

Add skip connection in latent space before/after processor.

convert_: str

Keep OmegaConf containers when instantiating — model code uses attribute-style access throughout.

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.models.schemas.models.TransportSourceConfig(*, kind: Literal['default', 'zero', 'gaussian', 'reference_state'] = 'default', scale: Annotated[float, Ge(ge=0)] = 1.0, noise_scale: Annotated[float, Ge(ge=0)] = 0.0)

Bases: BaseModel

kind: Literal['default', 'zero', 'gaussian', 'reference_state']

Starting field used before the transport objective moves toward the target.

scale: NonNegativeFloat

Multiplier applied to the starting/source field.

noise_scale: NonNegativeFloat

Additional additive Gaussian noise applied to the starting/source field.

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.models.schemas.models.TransportConfig(*, objective: Literal['edm_diffusion', 'stochastic_interpolant']='edm_diffusion', sigma_data: Annotated[float, ~annotated_types.Gt(gt=0)] = 1.0, noise_channels: Annotated[int, ~annotated_types.Gt(gt=0)] = 32, noise_cond_dim: Annotated[int, ~annotated_types.Gt(gt=0)] = 16, sigma_max: Annotated[float, ~annotated_types.Gt(gt=0)] = 100.0, sigma_min: Annotated[float, ~annotated_types.Gt(gt=0)] = 0.02, rho: Annotated[float, ~annotated_types.Gt(gt=0)] = 7.0, si_alpha_schedule: Literal['linear'] = 'linear', si_beta_schedule: Literal['linear', 'quadratic']='linear', si_sigma_schedule: Literal['brownian_bridge', 'quadratic_bridge']='brownian_bridge', source: TransportSourceConfig = <factory>, si_noise_scale: Annotated[float, ~annotated_types.Ge(ge=0)] = 1.0, training_condition: dict = <factory>, noise_embedder: dict = <factory>, inference_defaults: dict = <factory>)

Bases: BaseModel

objective: Literal['edm_diffusion', 'stochastic_interpolant']

Training and sampling objective used by the transport model.

sigma_data: PositiveFloat

Typical data scale used by EDM diffusion.

noise_channels: PositiveInt

Number of channels in the noise or bridge-time embedding.

noise_cond_dim: PositiveInt

Size of the conditioning vector passed to conditional layers.

sigma_max: PositiveFloat

Maximum EDM diffusion noise level used during training.

sigma_min: PositiveFloat

Minimum EDM diffusion noise level used during training.

rho: PositiveFloat

Shape parameter for the Karras EDM noise schedule.

si_alpha_schedule: Literal['linear']

Schedule for how strongly the SI bridge keeps the source field.

si_beta_schedule: Literal['linear', 'quadratic']

Schedule for how strongly the SI bridge moves toward the target field.

si_sigma_schedule: Literal['brownian_bridge', 'quadratic_bridge']

Schedule for the SI bridge-noise amplitude.

source: TransportSourceConfig

Configuration for the starting/source field.

si_noise_scale: NonNegativeFloat

Overall scale of the stochastic-interpolant bridge noise.

training_condition: dict

Distribution used to sample one training noise level or bridge time per sample.

noise_embedder: dict

Hydra configuration for embedding the current noise level or bridge time.

inference_defaults: dict

Default sampler parameters used during inference.

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.models.schemas.models.TransportModel(*, _target_: ~anemoi.models.schemas.models.DefinedModels, hidden_nodes_name: str | list[str], latent_skip: bool = True, _convert_: str = 'none', transport: ~anemoi.models.schemas.models.TransportConfig = <factory>)

Bases: Model

transport: TransportConfig

Transport model objective, path, conditioning, and inference configuration.

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.models.schemas.models.TrainableParameters(*, data: Annotated[int, Ge(ge=0)], hidden: Annotated[int, Ge(ge=0)])

Bases: BaseModel

data: NonNegativeInt

Size of the learnable data node tensor. Default to 8.

hidden: NonNegativeInt

Size of the learnable hidden node tensor. Default to 8.

model_config = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.models.schemas.models.ReluBoundingSchema(*, _target_: Literal['anemoi.models.layers.bounding.ReluBounding'], variables: list[str])

Bases: BaseModel

target_: Literal['anemoi.models.layers.bounding.ReluBounding']

Relu bounding object defined in anemoi.models.layers.bounding.

variables: list[str]

List of variables to bound using the Relu method.

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.models.schemas.models.LeakyReluBoundingSchema(*, _target_: Literal['anemoi.models.layers.bounding.LeakyReluBounding'], variables: list[str])

Bases: ReluBoundingSchema

target_: Literal['anemoi.models.layers.bounding.LeakyReluBounding']

Leaky Relu bounding object defined in anemoi.models.layers.bounding.

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.models.schemas.models.FractionBoundingSchema(*, _target_: Literal['anemoi.models.layers.bounding.FractionBounding'], variables: list[str], min_val: float, max_val: float, total_var: str)

Bases: BaseModel

target_: Literal['anemoi.models.layers.bounding.FractionBounding']

Fraction bounding object defined in anemoi.models.layers.bounding.

variables: list[str]

List of variables to bound using the hard tanh fraction method.

min_val: float

The minimum value for the HardTanh activation. Correspond to the minimum fraction of the total_var.

max_val: float

The maximum value for the HardTanh activation. Correspond to the maximum fraction of the total_var.

total_var: str

Variable from which the secondary variables are derived. For example, convective precipitation should be a fraction of total precipitation.

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.models.schemas.models.LeakyFractionBoundingSchema(*, _target_: Literal['anemoi.models.layers.bounding.LeakyFractionBounding'], variables: list[str], min_val: float, max_val: float, total_var: str)

Bases: FractionBoundingSchema

target_: Literal['anemoi.models.layers.bounding.LeakyFractionBounding']

Leaky fraction bounding object defined in anemoi.models.layers.bounding.

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.models.schemas.models.HardtanhBoundingSchema(*, _target_: Literal['anemoi.models.layers.bounding.HardtanhBounding'], variables: list[str], min_val: float, max_val: float)

Bases: BaseModel

target_: Literal['anemoi.models.layers.bounding.HardtanhBounding']

Hard tanh bounding method function from anemoi.models.layers.bounding.

variables: list[str]

List of variables to bound using the hard tanh method.

min_val: float

The minimum value for the HardTanh activation.

max_val: float

The maximum value for the HardTanh activation.

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.models.schemas.models.LeakyHardtanhBoundingSchema(*, _target_: Literal['anemoi.models.layers.bounding.LeakyHardtanhBounding'], variables: list[str], min_val: float, max_val: float)

Bases: HardtanhBoundingSchema

target_: Literal['anemoi.models.layers.bounding.LeakyHardtanhBounding']

Leaky hard tanh bounding method function from anemoi.models.layers.bounding.

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.models.schemas.models.NormalizedReluBoundingSchema(*, _target_: Literal['anemoi.models.layers.bounding.NormalizedReluBounding'], variables: list[str], min_val: list[float], normalizer: list[str])

Bases: BaseModel

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.models.schemas.models.NormalizedLeakyReluBoundingSchema(*, _target_: Literal['anemoi.models.layers.bounding.NormalizedLeakyReluBounding'], variables: list[str], min_val: list[float], normalizer: list[str])

Bases: NormalizedReluBoundingSchema

target_: Literal['anemoi.models.layers.bounding.NormalizedLeakyReluBounding']

Leaky normalized Relu bounding object defined in anemoi.models.layers.bounding.

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.models.schemas.models.NoOutputMaskSchema(*, _target_: Literal['anemoi.training.utils.masks.NoOutputMask'])

Bases: BaseModel

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.models.schemas.models.Boolean1DSchema(*, _target_: Literal['anemoi.training.utils.masks.Boolean1DMask'], attribute_name: str)

Bases: BaseModel

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.models.schemas.models.BaseModelSchema(*, num_channels: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], keep_batch_sharded: bool = True, model: ~anemoi.models.schemas.models.Model = <factory>, trainable_parameters: ~anemoi.models.schemas.models.TrainableParameters = <factory>, bounding: list[~typing.Annotated[~anemoi.models.schemas.models.ReluBoundingSchema | ~anemoi.models.schemas.models.LeakyReluBoundingSchema | ~anemoi.models.schemas.models.FractionBoundingSchema | ~anemoi.models.schemas.models.LeakyFractionBoundingSchema | ~anemoi.models.schemas.models.HardtanhBoundingSchema | ~anemoi.models.schemas.models.LeakyHardtanhBoundingSchema | ~anemoi.models.schemas.models.NormalizedReluBoundingSchema | ~anemoi.models.schemas.models.NormalizedLeakyReluBoundingSchema, FieldInfo(annotation=NoneType, required=True, discriminator='target_')]], output_mask: ~anemoi.models.schemas.models.NoOutputMaskSchema | ~anemoi.models.schemas.models.Boolean1DSchema, latent_skip: bool = True, processor: ~anemoi.models.schemas.processor.NoOpProcessorSchema | ~anemoi.models.schemas.processor.GNNProcessorSchema | ~anemoi.models.schemas.processor.GraphTransformerProcessorSchema | ~anemoi.models.schemas.processor.TransformerProcessorSchema | ~anemoi.models.schemas.processor.PointWiseMLPProcessorSchema, encoder: ~anemoi.models.schemas.encoder.GNNEncoderSchema | ~anemoi.models.schemas.encoder.GraphTransformerEncoderSchema | ~anemoi.models.schemas.encoder.TransformerEncoderSchema | ~anemoi.models.schemas.encoder.PointWiseForwardMapperSchema, decoder: ~anemoi.models.schemas.decoder.GNNDecoderSchema | ~anemoi.models.schemas.decoder.GraphTransformerDecoderSchema | ~anemoi.models.schemas.decoder.TransformerDecoderSchema | ~anemoi.models.schemas.decoder.PointWiseBackwardMapperSchema, residual: ~anemoi.models.schemas.residual.SkipConnectionSchema | ~anemoi.models.schemas.residual.TruncatedConnectionSchema | ~anemoi.models.schemas.residual.ScalarOrnsteinConnectionSchema | ~anemoi.models.schemas.residual.SpectralOrnsteinConnectionSchema, compile: list[dict[str, ~typing.Any]] | None = None)

Bases: BaseModel

num_channels: NonNegativeInt

Feature tensor size in the hidden space.

keep_batch_sharded: bool

Keep the input batch and the output of the model sharded

model_config = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model: Model

Model schema.

trainable_parameters: TrainableParameters

Learnable node and edge parameters.

bounding: list[Bounding]

List of bounding configuration applied in order to the specified variables.

output_mask: OutputMaskSchemas

Output mask

latent_skip: bool

Add skip connection in latent space before/after processor.

processor: NoOpProcessorSchema | GNNProcessorSchema | GraphTransformerProcessorSchema | TransformerProcessorSchema | PointWiseMLPProcessorSchema

GNN processor schema.

encoder: GNNEncoderSchema | GraphTransformerEncoderSchema | TransformerEncoderSchema | PointWiseForwardMapperSchema

GNN encoder schema.

residual: ResidualConnectionSchema

Residual connection schema.

compile: list[dict[str, Any]] | None

Modules to be compiled

class anemoi.models.schemas.models.NoOpNoiseInjectorSchema(*, _target_: Literal['anemoi.models.layers.ensemble.NoOpNoiseInjector'])

Bases: BaseModel

Schema for NoOpNoiseInjector - passes input through unchanged.

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

target_: Literal['anemoi.models.layers.ensemble.NoOpNoiseInjector']

No-op noise injector class

class anemoi.models.schemas.models.NoiseConditioningSchema(*, _target_: ~typing.Literal['anemoi.models.layers.ensemble.NoiseConditioning'], noise_std: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], noise_channels_dim: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], noise_mlp_hidden_dim: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], layer_kernels: dict[str, dict] | None = <factory>, noise_matrix: str | None = None, noise_edges_name: tuple[str, str, str] | None = None, edge_weight_attribute: str | None = None, row_normalize_noise_matrix: bool = False, autocast: bool = False)

Bases: BaseModel

Schema for NoiseConditioning - generates noise for conditioning.

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

target_: Literal['anemoi.models.layers.ensemble.NoiseConditioning']

Noise conditioning layer class

noise_std: NonNegativeInt

Standard deviation of the noise to be injected.

noise_channels_dim: NonNegativeInt

Number of channels in the noise tensor.

noise_mlp_hidden_dim: NonNegativeInt

Hidden dimension of the MLP used to process the noise.

layer_kernels: dict[str, dict] | None

Settings related to custom kernels for encoder processor and decoder blocks

noise_matrix: str | None

Path to the noise projection matrix file (.npz). If None, no projection is applied.

noise_edges_name: tuple[str, str, str] | None

Edge type identifier (src, relation, dst) for graph-based noise projection.

edge_weight_attribute: str | None

Optional edge attribute name for graph-based noise projection weights.

row_normalize_noise_matrix: bool

Whether to row-normalize the noise projection matrix weights.

autocast: bool

Whether to use autocast for the noise projection matrix operations.

class anemoi.models.schemas.models.NoiseInjectorSchema(*, _target_: ~typing.Literal['anemoi.models.layers.ensemble.NoiseInjector'], noise_std: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], noise_channels_dim: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], noise_mlp_hidden_dim: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], layer_kernels: dict[str, dict] | None = <factory>)

Bases: BaseModel

Schema for NoiseInjector - injects noise directly into input tensor.

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

target_: Literal['anemoi.models.layers.ensemble.NoiseInjector']

Noise injector layer class

noise_std: NonNegativeInt

Standard deviation of the noise to be injected.

noise_channels_dim: NonNegativeInt

Number of channels in the noise tensor.

noise_mlp_hidden_dim: NonNegativeInt

Hidden dimension of the MLP used to process the noise.

layer_kernels: dict[str, dict] | None

Settings related to custom kernels for encoder processor and decoder blocks

class anemoi.models.schemas.models.EnsModelSchema(*, num_channels: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], keep_batch_sharded: bool = True, model: ~anemoi.models.schemas.models.Model = <factory>, trainable_parameters: ~anemoi.models.schemas.models.TrainableParameters = <factory>, bounding: list[~typing.Annotated[~anemoi.models.schemas.models.ReluBoundingSchema | ~anemoi.models.schemas.models.LeakyReluBoundingSchema | ~anemoi.models.schemas.models.FractionBoundingSchema | ~anemoi.models.schemas.models.LeakyFractionBoundingSchema | ~anemoi.models.schemas.models.HardtanhBoundingSchema | ~anemoi.models.schemas.models.LeakyHardtanhBoundingSchema | ~anemoi.models.schemas.models.NormalizedReluBoundingSchema | ~anemoi.models.schemas.models.NormalizedLeakyReluBoundingSchema, FieldInfo(annotation=NoneType, required=True, discriminator='target_')]], output_mask: ~anemoi.models.schemas.models.NoOutputMaskSchema | ~anemoi.models.schemas.models.Boolean1DSchema, latent_skip: bool = True, processor: ~anemoi.models.schemas.processor.NoOpProcessorSchema | ~anemoi.models.schemas.processor.GNNProcessorSchema | ~anemoi.models.schemas.processor.GraphTransformerProcessorSchema | ~anemoi.models.schemas.processor.TransformerProcessorSchema | ~anemoi.models.schemas.processor.PointWiseMLPProcessorSchema, encoder: ~anemoi.models.schemas.encoder.GNNEncoderSchema | ~anemoi.models.schemas.encoder.GraphTransformerEncoderSchema | ~anemoi.models.schemas.encoder.TransformerEncoderSchema | ~anemoi.models.schemas.encoder.PointWiseForwardMapperSchema, decoder: ~anemoi.models.schemas.decoder.GNNDecoderSchema | ~anemoi.models.schemas.decoder.GraphTransformerDecoderSchema | ~anemoi.models.schemas.decoder.TransformerDecoderSchema | ~anemoi.models.schemas.decoder.PointWiseBackwardMapperSchema, residual: ~anemoi.models.schemas.residual.SkipConnectionSchema | ~anemoi.models.schemas.residual.TruncatedConnectionSchema | ~anemoi.models.schemas.residual.ScalarOrnsteinConnectionSchema | ~anemoi.models.schemas.residual.SpectralOrnsteinConnectionSchema, compile: list[dict[str, ~typing.Any]] | None = None, noise_injector: ~anemoi.models.schemas.models.NoOpNoiseInjectorSchema | ~anemoi.models.schemas.models.NoiseConditioningSchema | ~anemoi.models.schemas.models.NoiseInjectorSchema, condition_on_residual: bool = False)

Bases: BaseModelSchema

model_config = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

noise_injector: NoiseInjectorUnion

Noise injection configuration. Use NoOpNoiseInjector to disable, NoiseConditioning for conditioning, or NoiseInjector for direct injection.

condition_on_residual: bool

Whether to condition the noise injection on the residual connection.

class anemoi.models.schemas.models.TransportModelSchema(*, num_channels: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], keep_batch_sharded: bool = True, model: ~anemoi.models.schemas.models.TransportModel = <factory>, trainable_parameters: ~anemoi.models.schemas.models.TrainableParameters = <factory>, bounding: list[~typing.Annotated[~anemoi.models.schemas.models.ReluBoundingSchema | ~anemoi.models.schemas.models.LeakyReluBoundingSchema | ~anemoi.models.schemas.models.FractionBoundingSchema | ~anemoi.models.schemas.models.LeakyFractionBoundingSchema | ~anemoi.models.schemas.models.HardtanhBoundingSchema | ~anemoi.models.schemas.models.LeakyHardtanhBoundingSchema | ~anemoi.models.schemas.models.NormalizedReluBoundingSchema | ~anemoi.models.schemas.models.NormalizedLeakyReluBoundingSchema, FieldInfo(annotation=NoneType, required=True, discriminator='target_')]], output_mask: ~anemoi.models.schemas.models.NoOutputMaskSchema | ~anemoi.models.schemas.models.Boolean1DSchema, latent_skip: bool = True, processor: ~anemoi.models.schemas.processor.NoOpProcessorSchema | ~anemoi.models.schemas.processor.GNNProcessorSchema | ~anemoi.models.schemas.processor.GraphTransformerProcessorSchema | ~anemoi.models.schemas.processor.TransformerProcessorSchema | ~anemoi.models.schemas.processor.PointWiseMLPProcessorSchema, encoder: ~anemoi.models.schemas.encoder.GNNEncoderSchema | ~anemoi.models.schemas.encoder.GraphTransformerEncoderSchema | ~anemoi.models.schemas.encoder.TransformerEncoderSchema | ~anemoi.models.schemas.encoder.PointWiseForwardMapperSchema, decoder: ~anemoi.models.schemas.decoder.GNNDecoderSchema | ~anemoi.models.schemas.decoder.GraphTransformerDecoderSchema | ~anemoi.models.schemas.decoder.TransformerDecoderSchema | ~anemoi.models.schemas.decoder.PointWiseBackwardMapperSchema, residual: ~anemoi.models.schemas.residual.SkipConnectionSchema | ~anemoi.models.schemas.residual.TruncatedConnectionSchema | ~anemoi.models.schemas.residual.ScalarOrnsteinConnectionSchema | ~anemoi.models.schemas.residual.SpectralOrnsteinConnectionSchema, compile: list[dict[str, ~typing.Any]] | None = None)

Bases: BaseModelSchema

model_config = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model: TransportModel

Transport model schema.

class anemoi.models.schemas.models.TransportTendModelSchema(*, num_channels: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], keep_batch_sharded: bool = True, model: ~anemoi.models.schemas.models.TransportModel = <factory>, trainable_parameters: ~anemoi.models.schemas.models.TrainableParameters = <factory>, bounding: list[~typing.Annotated[~anemoi.models.schemas.models.ReluBoundingSchema | ~anemoi.models.schemas.models.LeakyReluBoundingSchema | ~anemoi.models.schemas.models.FractionBoundingSchema | ~anemoi.models.schemas.models.LeakyFractionBoundingSchema | ~anemoi.models.schemas.models.HardtanhBoundingSchema | ~anemoi.models.schemas.models.LeakyHardtanhBoundingSchema | ~anemoi.models.schemas.models.NormalizedReluBoundingSchema | ~anemoi.models.schemas.models.NormalizedLeakyReluBoundingSchema, FieldInfo(annotation=NoneType, required=True, discriminator='target_')]], output_mask: ~anemoi.models.schemas.models.NoOutputMaskSchema | ~anemoi.models.schemas.models.Boolean1DSchema, latent_skip: bool = True, processor: ~anemoi.models.schemas.processor.NoOpProcessorSchema | ~anemoi.models.schemas.processor.GNNProcessorSchema | ~anemoi.models.schemas.processor.GraphTransformerProcessorSchema | ~anemoi.models.schemas.processor.TransformerProcessorSchema | ~anemoi.models.schemas.processor.PointWiseMLPProcessorSchema, encoder: ~anemoi.models.schemas.encoder.GNNEncoderSchema | ~anemoi.models.schemas.encoder.GraphTransformerEncoderSchema | ~anemoi.models.schemas.encoder.TransformerEncoderSchema | ~anemoi.models.schemas.encoder.PointWiseForwardMapperSchema, decoder: ~anemoi.models.schemas.decoder.GNNDecoderSchema | ~anemoi.models.schemas.decoder.GraphTransformerDecoderSchema | ~anemoi.models.schemas.decoder.TransformerDecoderSchema | ~anemoi.models.schemas.decoder.PointWiseBackwardMapperSchema, residual: ~anemoi.models.schemas.residual.SkipConnectionSchema | ~anemoi.models.schemas.residual.TruncatedConnectionSchema | ~anemoi.models.schemas.residual.ScalarOrnsteinConnectionSchema | ~anemoi.models.schemas.residual.SpectralOrnsteinConnectionSchema, compile: list[dict[str, ~typing.Any]] | None = None, condition_on_residual: bool = False)

Bases: TransportModelSchema

model_config = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

condition_on_residual: bool

Whether to condition the noise injection on the residual connection.

class anemoi.models.schemas.models.HierarchicalModelSchema(*, num_channels: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], keep_batch_sharded: bool = True, model: ~anemoi.models.schemas.models.Model = <factory>, trainable_parameters: ~anemoi.models.schemas.models.TrainableParameters = <factory>, bounding: list[~typing.Annotated[~anemoi.models.schemas.models.ReluBoundingSchema | ~anemoi.models.schemas.models.LeakyReluBoundingSchema | ~anemoi.models.schemas.models.FractionBoundingSchema | ~anemoi.models.schemas.models.LeakyFractionBoundingSchema | ~anemoi.models.schemas.models.HardtanhBoundingSchema | ~anemoi.models.schemas.models.LeakyHardtanhBoundingSchema | ~anemoi.models.schemas.models.NormalizedReluBoundingSchema | ~anemoi.models.schemas.models.NormalizedLeakyReluBoundingSchema, FieldInfo(annotation=NoneType, required=True, discriminator='target_')]], output_mask: ~anemoi.models.schemas.models.NoOutputMaskSchema | ~anemoi.models.schemas.models.Boolean1DSchema, latent_skip: bool = True, processor: ~anemoi.models.schemas.processor.NoOpProcessorSchema | ~anemoi.models.schemas.processor.GNNProcessorSchema | ~anemoi.models.schemas.processor.GraphTransformerProcessorSchema | ~anemoi.models.schemas.processor.TransformerProcessorSchema | ~anemoi.models.schemas.processor.PointWiseMLPProcessorSchema, encoder: ~anemoi.models.schemas.encoder.GNNEncoderSchema | ~anemoi.models.schemas.encoder.GraphTransformerEncoderSchema | ~anemoi.models.schemas.encoder.TransformerEncoderSchema | ~anemoi.models.schemas.encoder.PointWiseForwardMapperSchema, decoder: ~anemoi.models.schemas.decoder.GNNDecoderSchema | ~anemoi.models.schemas.decoder.GraphTransformerDecoderSchema | ~anemoi.models.schemas.decoder.TransformerDecoderSchema | ~anemoi.models.schemas.decoder.PointWiseBackwardMapperSchema, residual: ~anemoi.models.schemas.residual.SkipConnectionSchema | ~anemoi.models.schemas.residual.TruncatedConnectionSchema | ~anemoi.models.schemas.residual.ScalarOrnsteinConnectionSchema | ~anemoi.models.schemas.residual.SpectralOrnsteinConnectionSchema, compile: list[dict[str, ~typing.Any]] | None = None, enable_hierarchical_level_processing: bool = False, level_process_num_layers: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)] = 1)

Bases: BaseModelSchema

model_config = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

enable_hierarchical_level_processing: bool

Toggle to do message passing at every downscaling and upscaling step

level_process_num_layers: NonNegativeInt

Number of message passing steps at each level

class anemoi.models.schemas.processor.NoOpProcessorSchema(*, _target_: Literal['anemoi.models.layers.processor.NoOpProcessor'])

Bases: BaseModel

target_: Literal['anemoi.models.layers.processor.NoOpProcessor']

No-op processor, used for ablations.

model_config = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.models.schemas.processor.GNNProcessorSchema(*, _convert_: str = 'all', trainable_size: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], num_chunks: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], cpu_offload: bool, gradient_checkpointing: bool = True, sub_graph_edge_attributes: list[str] = <factory>, mlp_extra_layers: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], mlp_hidden_ratio: ~typing.Annotated[float, ~annotated_types.Gt(gt=0)] = 1.0, mlp_implementation: ~typing.Literal['mlp', 'glu', 'swiglu', 'geglu', 'reglu'] = 'mlp', layer_kernels: dict[str, dict] | None = <factory>, _target_: ~typing.Literal['anemoi.models.layers.processor.GNNProcessor'], num_layers: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)])

Bases: GNNModelComponent

target_: Literal['anemoi.models.layers.processor.GNNProcessor']

GNN Processor object from anemoi.models.layers.processor.

num_layers: Annotated[int, Ge(ge=0)]

Number of layers of GNN processor. Default to 16.

num_chunks: Annotated[int, Ge(ge=0)]

Number of chunks to divide the layer into. Default to 2.

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.models.schemas.processor.GraphTransformerProcessorSchema(*, _convert_: str = 'all', cpu_offload: bool, gradient_checkpointing: bool = True, num_chunks: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], mlp_hidden_ratio: ~typing.Annotated[float, ~annotated_types.Gt(gt=0)], mlp_implementation: ~typing.Literal['mlp', 'glu', 'swiglu', 'geglu', 'reglu'] = 'mlp', num_heads: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], attn_channels: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)] | None = None, layer_kernels: dict[str, dict] | None = <factory>, _target_: ~typing.Literal['anemoi.models.layers.processor.GraphTransformerProcessor'], trainable_size: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], sub_graph_edge_attributes: list[str], num_layers: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], qk_norm: bool, **extra_data: ~typing.Any)

Bases: TransformerModelComponent

target_: Literal['anemoi.models.layers.processor.GraphTransformerProcessor']

Graph transformer processor object from anemoi.models.layers.processor.

trainable_size: Annotated[int, Ge(ge=0)]

Size of trainable parameters vector. Default to 8.

sub_graph_edge_attributes: list[str]

Edge attributes to consider in the processor features. Default [edge_length, endge_dirs].

num_layers: Annotated[int, Ge(ge=0)]

Number of layers of Graph Transformer processor. Default to 16.

num_chunks: Annotated[int, Ge(ge=0)]

Number of chunks to divide the layer into. Default to 2.

qk_norm: bool

Normalize the query and key vectors. Default to False.

model_config = {'extra': 'allow', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.models.schemas.processor.TransformerProcessorSchema(*, _convert_: str = 'all', cpu_offload: bool, gradient_checkpointing: bool = True, num_chunks: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], mlp_hidden_ratio: ~typing.Annotated[float, ~annotated_types.Gt(gt=0)], mlp_implementation: ~typing.Literal['mlp', 'glu', 'swiglu', 'geglu', 'reglu'] = 'mlp', num_heads: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], attn_channels: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)] | None = None, layer_kernels: dict[str, dict] | None = <factory>, _target_: ~typing.Literal['anemoi.models.layers.processor.TransformerProcessor'], num_layers: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], window_size: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)] | None, dropout_p: ~typing.Annotated[float, ~annotated_types.Ge(ge=0)], attention_implementation: str, qk_norm: bool, softcap: ~typing.Annotated[float, ~annotated_types.Ge(ge=0)], use_alibi_slopes: bool, **extra_data: ~typing.Any)

Bases: TransformerModelComponent

target_: Literal['anemoi.models.layers.processor.TransformerProcessor']

Transformer processor object from anemoi.models.layers.processor.

num_layers: Annotated[int, Ge(ge=0)]

Number of layers of Transformer processor. Default to 16.

model_config = {'extra': 'allow', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

num_chunks: Annotated[int, Ge(ge=0)]

Number of chunks to divide the layer into. Default to 2.

window_size: Annotated[int, Ge(ge=0)] | None

Attention window size along the longitude axis. Default to 512.

dropout_p: Annotated[float, Ge(ge=0)]

Dropout probability used for multi-head self attention, default 0.0

attention_implementation: str

Attention implementation to use. Default to ‘flash_attention’.

qk_norm: bool

Normalize the query and key vectors. Default to False.

softcap: Annotated[float, Ge(ge=0)]

Softcap value for attention. Default to 0.0.

use_alibi_slopes: bool

Use alibi slopes for attention implementation. Default to False.

class anemoi.models.schemas.processor.PointWiseMLPProcessorSchema(*, _convert_: str = 'all', num_chunks: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], cpu_offload: bool, gradient_checkpointing: bool = True, layer_kernels: dict[str, dict] | None = <factory>, _target_: ~typing.Literal['anemoi.models.layers.processor.PointWiseMLPProcessor'], num_layers: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], mlp_hidden_ratio: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], dropout_p: ~typing.Annotated[float, ~annotated_types.Ge(ge=0)] = 0.0)

Bases: PointWiseModelComponent

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

target_: Literal['anemoi.models.layers.processor.PointWiseMLPProcessor']

Transformer processor object from anemoi.models.layers.processor.

num_layers: Annotated[int, Ge(ge=0)]

Number of layers of Transformer processor.

mlp_hidden_ratio: Annotated[int, Ge(ge=0)]

Ratio of the hidden dimension to the processor channel dimension.

dropout_p: Annotated[float, Ge(ge=0)]

Dropout probability, default 0.0

class anemoi.models.schemas.encoder.GNNEncoderSchema(*, _convert_: str = 'all', trainable_size: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], num_chunks: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], cpu_offload: bool, gradient_checkpointing: bool = True, sub_graph_edge_attributes: list[str] = <factory>, mlp_extra_layers: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], mlp_hidden_ratio: ~typing.Annotated[float, ~annotated_types.Gt(gt=0)] = 1.0, mlp_implementation: ~typing.Literal['mlp', 'glu', 'swiglu', 'geglu', 'reglu'] = 'mlp', layer_kernels: dict[str, dict] | None = <factory>, _target_: ~typing.Literal['anemoi.models.layers.mapper.GNNForwardMapper'])

Bases: GNNModelComponent

target_: Literal['anemoi.models.layers.mapper.GNNForwardMapper']

GNN encoder object from anemoi.models.layers.mapper.

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.models.schemas.encoder.GraphTransformerEncoderSchema(*, _convert_: str = 'all', cpu_offload: bool, gradient_checkpointing: bool = True, num_chunks: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], mlp_hidden_ratio: ~typing.Annotated[float, ~annotated_types.Gt(gt=0)], mlp_implementation: ~typing.Literal['mlp', 'glu', 'swiglu', 'geglu', 'reglu'] = 'mlp', num_heads: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], attn_channels: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)] | None = None, layer_kernels: dict[str, dict] | None = <factory>, _target_: ~typing.Literal['anemoi.models.layers.mapper.GraphTransformerForwardMapper'], trainable_size: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], sub_graph_edge_attributes: list[str], qk_norm: bool, **extra_data: ~typing.Any)

Bases: TransformerModelComponent

target_: Literal['anemoi.models.layers.mapper.GraphTransformerForwardMapper']

Graph Transfromer Encoder object from anemoi.models.layers.mapper.

trainable_size: Annotated[int, Ge(ge=0)]

Size of trainable parameters vector. Default to 8.

sub_graph_edge_attributes: list[str]

Edge attributes to consider in the encoder features.

qk_norm: bool

Normalize the query and key vectors. Default to False.

model_config = {'extra': 'allow', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.models.schemas.encoder.TransformerEncoderSchema(*, _convert_: str = 'all', cpu_offload: bool, gradient_checkpointing: bool = True, num_chunks: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], mlp_hidden_ratio: ~typing.Annotated[float, ~annotated_types.Gt(gt=0)], mlp_implementation: ~typing.Literal['mlp', 'glu', 'swiglu', 'geglu', 'reglu'] = 'mlp', num_heads: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], attn_channels: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)] | None = None, layer_kernels: dict[str, dict] | None = <factory>, _target_: ~typing.Literal['anemoi.models.layers.mapper.TransformerForwardMapper'], window_size: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)] | None, dropout_p: ~typing.Annotated[float, ~annotated_types.Ge(ge=0)], attention_implementation: str, softcap: ~typing.Annotated[float, ~annotated_types.Ge(ge=0)], use_alibi_slopes: bool, **extra_data: ~typing.Any)

Bases: TransformerModelComponent

target_: Literal['anemoi.models.layers.mapper.TransformerForwardMapper']

Transformer Encoder object from anemoi.models.layers.mapper.

model_config = {'extra': 'allow', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

window_size: Annotated[int, Ge(ge=0)] | None

Attention window size along the longitude axis. Default to 512.

dropout_p: Annotated[float, Ge(ge=0)]

Dropout probability used for multi-head self attention, default 0.0

attention_implementation: str

Attention implementation to use. Default to ‘flash_attention’.

softcap: Annotated[float, Ge(ge=0)]

Softcap value for attention. Default to 0.0.

use_alibi_slopes: bool

Use alibi slopes for attention implementation. Default to False.

class anemoi.models.schemas.encoder.PointWiseForwardMapperSchema(*, _convert_: str = 'all', cpu_offload: bool, gradient_checkpointing: bool = True, layer_kernels: dict[str, dict] | None = <factory>, _target_: ~typing.Literal['anemoi.models.layers.mapper.PointWiseForwardMapper'])

Bases: PointWiseMapperComponent

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

target_: Literal['anemoi.models.layers.mapper.PointWiseForwardMapper']

Point-wise encoder object from anemoi.models.layers.mapper.

class anemoi.models.schemas.decoder.GNNDecoderSchema(*, _convert_: str = 'all', trainable_size: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], num_chunks: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], cpu_offload: bool, gradient_checkpointing: bool = True, sub_graph_edge_attributes: list[str] = <factory>, mlp_extra_layers: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], mlp_hidden_ratio: ~typing.Annotated[float, ~annotated_types.Gt(gt=0)] = 1.0, mlp_implementation: ~typing.Literal['mlp', 'glu', 'swiglu', 'geglu', 'reglu'] = 'mlp', layer_kernels: dict[str, dict] | None = <factory>, _target_: ~typing.Literal['anemoi.models.layers.mapper.GNNBackwardMapper'])

Bases: GNNModelComponent

target_: Literal['anemoi.models.layers.mapper.GNNBackwardMapper']

GNN decoder object from anemoi.models.layers.mapper.

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.models.schemas.decoder.GraphTransformerDecoderSchema(*, _convert_: str = 'all', cpu_offload: bool, gradient_checkpointing: bool = True, num_chunks: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], mlp_hidden_ratio: ~typing.Annotated[float, ~annotated_types.Gt(gt=0)], mlp_implementation: ~typing.Literal['mlp', 'glu', 'swiglu', 'geglu', 'reglu'] = 'mlp', num_heads: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], attn_channels: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)] | None = None, layer_kernels: dict[str, dict] | None = <factory>, _target_: ~typing.Literal['anemoi.models.layers.mapper.GraphTransformerBackwardMapper'], trainable_size: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], sub_graph_edge_attributes: list[str], qk_norm: bool, initialise_data_extractor_zero: bool, **extra_data: ~typing.Any)

Bases: TransformerModelComponent

target_: Literal['anemoi.models.layers.mapper.GraphTransformerBackwardMapper']

Graph Transformer Decoder object from anemoi.models.layers.mapper.

trainable_size: Annotated[int, Ge(ge=0)]

Size of trainable parameters vector. Default to 8.

sub_graph_edge_attributes: list[str]

Edge attributes to consider in the decoder features. Default to [edge_length, edge_dirs]

qk_norm: bool

Normalize the query and key vectors. Default to False.

initialise_data_extractor_zero: bool

Initialise the data extractor with zeros. Default to False.

model_config = {'extra': 'allow', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.models.schemas.decoder.TransformerDecoderSchema(*, _convert_: str = 'all', cpu_offload: bool, gradient_checkpointing: bool = True, num_chunks: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], mlp_hidden_ratio: ~typing.Annotated[float, ~annotated_types.Gt(gt=0)], mlp_implementation: ~typing.Literal['mlp', 'glu', 'swiglu', 'geglu', 'reglu'] = 'mlp', num_heads: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], attn_channels: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)] | None = None, layer_kernels: dict[str, dict] | None = <factory>, _target_: ~typing.Literal['anemoi.models.layers.mapper.TransformerBackwardMapper'], window_size: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)] | None, dropout_p: ~typing.Annotated[float, ~annotated_types.Ge(ge=0)], attention_implementation: str, softcap: ~typing.Annotated[float, ~annotated_types.Ge(ge=0)], use_alibi_slopes: bool, use_rotary_embeddings: bool, **extra_data: ~typing.Any)

Bases: TransformerModelComponent

target_: Literal['anemoi.models.layers.mapper.TransformerBackwardMapper']

Transformer Encoder object from anemoi.models.layers.mapper.

window_size: Annotated[int, Ge(ge=0)] | None

Attention window size along the longitude axis. Default to 512.

model_config = {'extra': 'allow', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

dropout_p: Annotated[float, Ge(ge=0)]

Dropout probability used for multi-head self attention, default 0.0

attention_implementation: str

Attention implementation to use. Default to ‘flash_attention’.

softcap: Annotated[float, Ge(ge=0)]

Softcap value for attention. Default to 0.0.

use_alibi_slopes: bool

Use alibi slopes for attention implementation. Default to False.

use_rotary_embeddings: bool

Use rotary embeddings for attention implementation. Default to False.

class anemoi.models.schemas.decoder.PointWiseBackwardMapperSchema(*, _convert_: str = 'all', cpu_offload: bool, gradient_checkpointing: bool = True, layer_kernels: dict[str, dict] | None = <factory>, _target_: ~typing.Literal['anemoi.models.layers.mapper.PointWiseBackwardMapper'], initialise_data_extractor_zero: bool = False)

Bases: PointWiseMapperComponent

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

target_: Literal['anemoi.models.layers.mapper.PointWiseBackwardMapper']

Point-wise decoder object from anemoi.models.layers.mapper.

initialise_data_extractor_zero: bool

Initialise the data extractor with zeros. Default to False.