hf-internal-testing/tiny-random-WeatherNext2ForWeatherForecasting

Tiny random WeatherNext 2

A randomly initialized WeatherNext2ForWeatherForecasting for use in the Transformers test suite. 13K parameters, 152 KB. It forecasts nothing: the weights are noise and the feature extractor's statistics are zero mean and unit variance.

What it is for is the geometry. WeatherNext 2 has no learned positional encodings, so the mesh, both grid-mesh graphs and the banded attention mask are built once at conversion time, with scipy and trimesh, and stored in the checkpoint as buffers. A model built from a config alone gets a placeholder instead, which is enough to run but is not a mesh. This checkpoint carries the real thing at a size CI can download, so the tests that depend on a real mesh have one.

The grid is 15 degrees and the mesh is a twice-split icosahedron: 162 mesh nodes over 312 grid points, an attention bandwidth of 43, so four attention blocks. Four is the point, since it gives an interior block with a neighbour on either side.

Regenerating it

import torch
from transformers import WeatherNext2Config, WeatherNext2FeatureExtractor, WeatherNext2ForWeatherForecasting
from transformers.models.weathernext2.convert_weathernext2_original_checkpoint import geometry_state_dict

config = WeatherNext2Config(
    hidden_size=16, intermediate_size=32, num_hidden_layers=2, num_attention_heads=2,
    edge_hidden_size=4, noise_channels=4, mesh_splits=2, attention_k_hop=2,
    grid_latitudes=13, grid_longitudes=24, pressure_levels=(500, 850), aggregate_normalization=None,
)
geometry = geometry_state_dict(config)  # needs scipy and trimesh; sets the two derived sizes

torch.manual_seed(0)
model = WeatherNext2ForWeatherForecasting(config)
model.load_state_dict(geometry, strict=False)

The feature extractor is built from the same config with zero means and unit standard deviations for every variable.