From 2befa14e5c8cbd44740b89a730848da2d444b77d Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Mon, 10 Aug 2026 06:14:03 +0000
Subject: [PATCH 01/18] Initial plan


From ad8414048156d9880ae63b57c2692657d6252c80 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Mon, 10 Aug 2026 06:22:48 +0000
Subject: [PATCH 02/18] Add WeatherNext ONNX demo

Signed-off-by: GitHub <noreply@github.com>

Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
---
 examples/weathernext.py           | 314 ++++++++++++++++++++++++++++++
 tests/weathernext_example_test.py |  47 +++++
 2 files changed, 361 insertions(+)
 create mode 100644 examples/weathernext.py
 create mode 100644 tests/weathernext_example_test.py

diff --git a/examples/weathernext.py b/examples/weathernext.py
new file mode 100644
index 000000000..456582cdd
--- /dev/null
+++ b/examples/weathernext.py
@@ -0,0 +1,314 @@
+#!/usr/bin/env python
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""WeatherNext-style ONNX conversion demo.
+
+WeatherNext 2 is a JAX/Haiku + xarray model rather than a HuggingFace
+``transformers`` model, so it does not fit the normal ``mobius build`` path.
+This example demonstrates the ONNX workflow for that family of models by
+defining the same high-level one-step forecast contract used by WeatherNext:
+
+``input weather grid + forcings + stochastic noise → next weather grid``.
+
+The tiny model below uses WeatherNext-like graph data flow:
+
+1. encode lat/lon grid variables at each grid cell,
+2. aggregate grid cells onto a mesh,
+3. update the mesh latent state,
+4. project mesh latents back to the grid, and
+5. decode per-grid-cell forecast variables.
+
+It intentionally uses small deterministic weights so the demo can be run
+without downloading WeatherNext checkpoints.  The graph I/O and component
+boundaries are the pieces to keep when replacing the toy modules with a full
+translation of google-deepmind/weathernext's Haiku modules and ``.npz``
+checkpoints.
+
+Usage::
+
+    python examples/weathernext.py output/weathernext-mini --validate
+
+    python examples/weathernext.py output/weathernext-mini \
+        --lat 8 --lon 16 --mesh-nodes 12 --hidden-size 32
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+from dataclasses import dataclass
+
+import numpy as np
+import onnx_ir as ir
+from onnxscript import OpBuilder, nn
+
+from mobius import ArchitectureConfig, ModelPackage, build_from_module
+from mobius.tasks import ModelTask
+from mobius.tasks._base import _make_graph, _make_model
+
+
+@dataclass(frozen=True)
+class WeatherNextDemoShape:
+    """Concrete shape for the one-step WeatherNext forecast demo."""
+
+    lat: int
+    lon: int
+    mesh_nodes: int
+    input_variables: int
+    forcing_variables: int
+    noise_channels: int
+    output_variables: int
+    hidden_size: int
+
+    @property
+    def grid_points(self) -> int:
+        return self.lat * self.lon
+
+    @property
+    def encoder_channels(self) -> int:
+        return self.input_variables + self.forcing_variables + self.noise_channels
+
+
+def _make_parameter(rng: np.random.Generator, shape: tuple[int, ...]) -> nn.Parameter:
+    values = rng.standard_normal(shape).astype(np.float32) * 0.05
+    return nn.Parameter(list(shape), data=ir.tensor(values))
+
+
+class DemoLinear(nn.Module):
+    """Small deterministic ``Linear`` layer for a self-contained runnable demo."""
+
+    def __init__(self, rng: np.random.Generator, in_features: int, out_features: int):
+        super().__init__()
+        self.weight = _make_parameter(rng, (out_features, in_features))
+        self.bias = _make_parameter(rng, (out_features,))
+
+    def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value:
+        # Project the trailing feature dimension: [..., in_features] -> [..., out_features].
+        x = op.MatMul(x, op.Transpose(self.weight, perm=[1, 0]))
+        return op.Add(x, self.bias)
+
+
+class WeatherNextGridMeshBlock(nn.Module):
+    """One grid→mesh→grid block mirroring WeatherNext's graph-forecast data flow."""
+
+    def __init__(self, rng: np.random.Generator, shape: WeatherNextDemoShape):
+        super().__init__()
+        self._shape = shape
+        self.grid_encoder = DemoLinear(rng, shape.encoder_channels, shape.hidden_size)
+        self.mesh_update_in = DemoLinear(rng, shape.hidden_size, 4 * shape.hidden_size)
+        self.mesh_update_out = DemoLinear(rng, 4 * shape.hidden_size, shape.hidden_size)
+        self.grid_decoder = DemoLinear(rng, shape.hidden_size, shape.output_variables)
+
+        grid_to_mesh = _projection_matrix(shape.mesh_nodes, shape.grid_points)
+        mesh_to_grid = _projection_matrix(shape.grid_points, shape.mesh_nodes)
+        self.grid_to_mesh = nn.Parameter(
+            [shape.mesh_nodes, shape.grid_points], data=ir.tensor(grid_to_mesh)
+        )
+        self.mesh_to_grid = nn.Parameter(
+            [shape.grid_points, shape.mesh_nodes], data=ir.tensor(mesh_to_grid)
+        )
+
+    def forward(
+        self,
+        op: OpBuilder,
+        input_state: ir.Value,
+        forcings: ir.Value,
+        sample_noise: ir.Value,
+    ) -> ir.Value:
+        s = self._shape
+
+        # Concatenate per-cell weather variables, known future forcings, and FGN noise:
+        # [B, lat, lon, input+forcing+noise].
+        grid_features = op.Concat(input_state, forcings, sample_noise, axis=-1)
+
+        # Encode each lat/lon cell independently, then flatten the grid to points:
+        # [B, lat, lon, hidden] -> [B, grid_points, hidden].
+        grid_latent = op.Tanh(self.grid_encoder(op, grid_features))
+        grid_points = op.Reshape(grid_latent, [0, s.grid_points, s.hidden_size])
+
+        # Aggregate grid points onto the mesh with a fixed sparse-style projection:
+        # [mesh_points, grid_points] @ [B, grid_points, hidden] -> [B, mesh_points, hidden].
+        mesh_latent = op.MatMul(self.grid_to_mesh, grid_points)
+
+        # A compact MLP stands in for WeatherNext's mesh GNN/update blocks.
+        mesh_delta = op.Tanh(self.mesh_update_in(op, mesh_latent))
+        mesh_latent = op.Add(mesh_latent, self.mesh_update_out(op, mesh_delta))
+
+        # Decode mesh latents back onto the lat/lon grid and add the encoded-grid residual:
+        # [grid_points, mesh_points] @ [B, mesh_points, hidden] -> [B, grid_points, hidden].
+        grid_delta = op.MatMul(self.mesh_to_grid, mesh_latent)
+        grid_points = op.Add(grid_points, grid_delta)
+
+        # Return a one-step forecast grid: [B, lat, lon, output_variables].
+        forecast_points = self.grid_decoder(op, grid_points)
+        return op.Reshape(forecast_points, [0, s.lat, s.lon, s.output_variables])
+
+
+class WeatherNextDemoTask(ModelTask):
+    """Task wiring for a one-step WeatherNext-style forecast graph."""
+
+    model_roles = {"model": "encoder"}
+
+    def __init__(self, shape: WeatherNextDemoShape):
+        self._shape = shape
+
+    def build(self, module: nn.Module, config: ArchitectureConfig) -> ModelPackage:
+        batch = ir.SymbolicDim("batch")
+        s = self._shape
+
+        graph, builder = _make_graph("weathernext_one_step_forecast")
+        op = builder.op
+
+        input_state = builder.input(
+            "input_state",
+            dtype=config.dtype,
+            shape=[batch, s.lat, s.lon, s.input_variables],
+        )
+        forcings = builder.input(
+            "forcings",
+            dtype=config.dtype,
+            shape=[batch, s.lat, s.lon, s.forcing_variables],
+        )
+        sample_noise = builder.input(
+            "sample_noise",
+            dtype=config.dtype,
+            shape=[batch, s.lat, s.lon, s.noise_channels],
+        )
+
+        next_state = module(op, input_state, forcings, sample_noise)
+        builder.add_output(next_state, "next_state")
+
+        return ModelPackage({"model": _make_model(graph)}, config=config)
+
+
+def _projection_matrix(rows: int, cols: int) -> np.ndarray:
+    """Create deterministic normalized projections between grid and mesh points."""
+
+    row_positions = np.linspace(0.0, 1.0, rows, dtype=np.float32)[:, None]
+    col_positions = np.linspace(0.0, 1.0, cols, dtype=np.float32)[None, :]
+    distance = np.abs(row_positions - col_positions)
+    weights = np.maximum(1.0 - 2.0 * distance, 0.0)
+    weights += 1e-3
+    weights /= weights.sum(axis=1, keepdims=True)
+    return weights.astype(np.float32)
+
+
+def build_weathernext_demo_package(shape: WeatherNextDemoShape, dtype: ir.DataType) -> ModelPackage:
+    """Build the demo WeatherNext-style ONNX package."""
+
+    rng = np.random.default_rng(20260810)
+    module = WeatherNextGridMeshBlock(rng, shape)
+    config = ArchitectureConfig(
+        vocab_size=0,
+        hidden_size=shape.hidden_size,
+        intermediate_size=4 * shape.hidden_size,
+        num_hidden_layers=1,
+        num_attention_heads=1,
+        num_key_value_heads=1,
+        head_dim=shape.hidden_size,
+        dtype=dtype,
+    )
+    return build_from_module(module, config, task=WeatherNextDemoTask(shape))
+
+
+def _resolve_dtype(name: str) -> ir.DataType:
+    if name == "f32":
+        return ir.DataType.FLOAT
+    if name == "f16":
+        return ir.DataType.FLOAT16
+    raise ValueError(f"Unsupported dtype: {name}")
+
+
+def _validate_with_ort(output_dir: str, shape: WeatherNextDemoShape) -> None:
+    try:
+        import onnxruntime as ort
+    except ImportError:
+        print("onnxruntime is not installed; skipping validation.", file=sys.stderr)
+        return
+
+    model_path = os.path.join(output_dir, "model.onnx")
+    sess = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
+    rng = np.random.default_rng(42)
+    feeds = {
+        "input_state": rng.standard_normal(
+            (1, shape.lat, shape.lon, shape.input_variables), dtype=np.float32
+        ),
+        "forcings": rng.standard_normal(
+            (1, shape.lat, shape.lon, shape.forcing_variables), dtype=np.float32
+        ),
+        "sample_noise": rng.standard_normal(
+            (1, shape.lat, shape.lon, shape.noise_channels), dtype=np.float32
+        ),
+    }
+    (next_state,) = sess.run(None, feeds)
+    print(f"Validation output next_state shape: {next_state.shape}")
+    print(f"Validation output range: [{next_state.min():.6f}, {next_state.max():.6f}]")
+
+
+def _parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Build a runnable WeatherNext-style grid→mesh→grid ONNX demo.",
+    )
+    parser.add_argument("output_dir", help="Directory to save model.onnx and model.onnx.data.")
+    parser.add_argument("--lat", type=int, default=4, help="Number of latitude points.")
+    parser.add_argument("--lon", type=int, default=8, help="Number of longitude points.")
+    parser.add_argument("--mesh-nodes", type=int, default=6, help="Number of demo mesh nodes.")
+    parser.add_argument("--input-variables", type=int, default=5, help="Input weather channels.")
+    parser.add_argument("--forcing-variables", type=int, default=2, help="Known forcing channels.")
+    parser.add_argument("--noise-channels", type=int, default=2, help="FGN stochastic noise channels.")
+    parser.add_argument("--output-variables", type=int, default=5, help="Forecast weather channels.")
+    parser.add_argument("--hidden-size", type=int, default=16, help="Latent feature size.")
+    parser.add_argument("--dtype", choices=["f32", "f16"], default="f32", help="ONNX weight dtype.")
+    parser.add_argument("--validate", action="store_true", help="Run one ONNX Runtime inference.")
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = _parse_args()
+    shape = WeatherNextDemoShape(
+        lat=args.lat,
+        lon=args.lon,
+        mesh_nodes=args.mesh_nodes,
+        input_variables=args.input_variables,
+        forcing_variables=args.forcing_variables,
+        noise_channels=args.noise_channels,
+        output_variables=args.output_variables,
+        hidden_size=args.hidden_size,
+    )
+
+    if min(
+        shape.lat,
+        shape.lon,
+        shape.mesh_nodes,
+        shape.input_variables,
+        shape.forcing_variables,
+        shape.noise_channels,
+        shape.output_variables,
+        shape.hidden_size,
+    ) <= 0:
+        raise ValueError("All shape arguments must be positive.")
+
+    print("Building WeatherNext-style one-step forecast ONNX graph...")
+    print(f"  grid: {shape.lat} x {shape.lon} ({shape.grid_points} cells)")
+    print(f"  mesh nodes: {shape.mesh_nodes}")
+    print(f"  channels: input={shape.input_variables}, forcing={shape.forcing_variables}, "
+          f"noise={shape.noise_channels}, output={shape.output_variables}")
+
+    pkg = build_weathernext_demo_package(shape, dtype=_resolve_dtype(args.dtype))
+    model = pkg["model"]
+    num_nodes = model.graph.num_nodes
+    if callable(num_nodes):
+        num_nodes = num_nodes()
+    print(f"Built model with {num_nodes} ONNX nodes.")
+
+    pkg.save(args.output_dir, check_weights=True, progress_bar=False)
+    print(f"Saved WeatherNext demo package to {args.output_dir!r}.")
+
+    if args.validate:
+        _validate_with_ort(args.output_dir, shape)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/weathernext_example_test.py b/tests/weathernext_example_test.py
new file mode 100644
index 000000000..140f8d257
--- /dev/null
+++ b/tests/weathernext_example_test.py
@@ -0,0 +1,47 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from __future__ import annotations
+
+import importlib.util
+import sys
+from pathlib import Path
+
+import onnx_ir as ir
+
+
+def _load_weathernext_example():
+    path = Path(__file__).parents[1] / "examples" / "weathernext.py"
+    spec = importlib.util.spec_from_file_location("weathernext_example", path)
+    assert spec is not None
+    assert spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def test_weathernext_demo_exports_one_step_forecast_model(tmp_path):
+    example = _load_weathernext_example()
+    shape = example.WeatherNextDemoShape(
+        lat=3,
+        lon=4,
+        mesh_nodes=5,
+        input_variables=2,
+        forcing_variables=1,
+        noise_channels=1,
+        output_variables=2,
+        hidden_size=8,
+    )
+
+    pkg = example.build_weathernext_demo_package(shape, dtype=ir.DataType.FLOAT)
+    pkg.save(tmp_path, check_weights=True, progress_bar=False)
+
+    model = ir.load(tmp_path / "model.onnx")
+    assert [value.name for value in model.graph.inputs] == [
+        "input_state",
+        "forcings",
+        "sample_noise",
+    ]
+    assert [value.name for value in model.graph.outputs] == ["next_state"]
+    assert model.graph.outputs[0].shape == ir.Shape(["batch", 3, 4, 2])

From d9a3806478696519d6704349770e5680b377ef3c Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Mon, 10 Aug 2026 06:25:24 +0000
Subject: [PATCH 03/18] Address WeatherNext demo review feedback

Signed-off-by: GitHub <noreply@github.com>

Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
---
 examples/weathernext.py | 33 ++++++++++++++++++++++++---------
 1 file changed, 24 insertions(+), 9 deletions(-)

diff --git a/examples/weathernext.py b/examples/weathernext.py
index 456582cdd..96bdbd3fb 100644
--- a/examples/weathernext.py
+++ b/examples/weathernext.py
@@ -42,11 +42,11 @@
 
 import numpy as np
 import onnx_ir as ir
-from onnxscript import OpBuilder, nn
+from onnxscript import GraphBuilder, OpBuilder, nn
 
-from mobius import ArchitectureConfig, ModelPackage, build_from_module
+import mobius
+from mobius import OPSET_VERSION, ArchitectureConfig, ModelPackage, build_from_module
 from mobius.tasks import ModelTask
-from mobius.tasks._base import _make_graph, _make_model
 
 
 @dataclass(frozen=True)
@@ -158,7 +158,7 @@ def build(self, module: nn.Module, config: ArchitectureConfig) -> ModelPackage:
         batch = ir.SymbolicDim("batch")
         s = self._shape
 
-        graph, builder = _make_graph("weathernext_one_step_forecast")
+        graph, builder = _make_demo_graph("weathernext_one_step_forecast")
         op = builder.op
 
         input_state = builder.input(
@@ -180,7 +180,25 @@ def build(self, module: nn.Module, config: ArchitectureConfig) -> ModelPackage:
         next_state = module(op, input_state, forcings, sample_noise)
         builder.add_output(next_state, "next_state")
 
-        return ModelPackage({"model": _make_model(graph)}, config=config)
+        return ModelPackage({"model": _make_demo_model(graph)}, config=config)
+
+
+def _make_demo_graph(name: str) -> tuple[ir.Graph, GraphBuilder]:
+    graph = ir.Graph(
+        [],
+        [],
+        nodes=[],
+        name=name,
+        opset_imports={"": OPSET_VERSION, "com.microsoft": 1},
+    )
+    return graph, GraphBuilder(graph)
+
+
+def _make_demo_model(graph: ir.Graph) -> ir.Model:
+    model = ir.Model(graph, ir_version=11)
+    model.producer_name = "mobius"
+    model.producer_version = mobius.__version__
+    return model
 
 
 def _projection_matrix(rows: int, cols: int) -> np.ndarray:
@@ -298,10 +316,7 @@ def main() -> None:
 
     pkg = build_weathernext_demo_package(shape, dtype=_resolve_dtype(args.dtype))
     model = pkg["model"]
-    num_nodes = model.graph.num_nodes
-    if callable(num_nodes):
-        num_nodes = num_nodes()
-    print(f"Built model with {num_nodes} ONNX nodes.")
+    print(f"Built model with {model.graph.num_nodes()} ONNX nodes.")
 
     pkg.save(args.output_dir, check_weights=True, progress_bar=False)
     print(f"Saved WeatherNext demo package to {args.output_dir!r}.")

From d58c232fc66edfdb43e80926e21f508573636149 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Mon, 10 Aug 2026 06:26:58 +0000
Subject: [PATCH 04/18] Refine WeatherNext demo graph wiring

Signed-off-by: GitHub <noreply@github.com>

Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
---
 examples/weathernext.py | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/examples/weathernext.py b/examples/weathernext.py
index 96bdbd3fb..082b9081d 100644
--- a/examples/weathernext.py
+++ b/examples/weathernext.py
@@ -126,7 +126,10 @@ def forward(
         # Encode each lat/lon cell independently, then flatten the grid to points:
         # [B, lat, lon, hidden] -> [B, grid_points, hidden].
         grid_latent = op.Tanh(self.grid_encoder(op, grid_features))
-        grid_points = op.Reshape(grid_latent, [0, s.grid_points, s.hidden_size])
+        grid_points = op.Reshape(
+            grid_latent,
+            op.Constant(value_ints=[0, s.grid_points, s.hidden_size]),
+        )
 
         # Aggregate grid points onto the mesh with a fixed sparse-style projection:
         # [mesh_points, grid_points] @ [B, grid_points, hidden] -> [B, mesh_points, hidden].
@@ -143,13 +146,16 @@ def forward(
 
         # Return a one-step forecast grid: [B, lat, lon, output_variables].
         forecast_points = self.grid_decoder(op, grid_points)
-        return op.Reshape(forecast_points, [0, s.lat, s.lon, s.output_variables])
+        return op.Reshape(
+            forecast_points,
+            op.Constant(value_ints=[0, s.lat, s.lon, s.output_variables]),
+        )
 
 
 class WeatherNextDemoTask(ModelTask):
     """Task wiring for a one-step WeatherNext-style forecast graph."""
 
-    model_roles = {"model": "encoder"}
+    model_roles = {"model": "forecast"}
 
     def __init__(self, shape: WeatherNextDemoShape):
         self._shape = shape

From 4fde44f280dee2002e5ff8a71e1f932f6d5167a3 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Mon, 10 Aug 2026 06:28:05 +0000
Subject: [PATCH 05/18] Use dynamic reshape in WeatherNext demo

Signed-off-by: GitHub <noreply@github.com>

Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
---
 examples/weathernext.py | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/examples/weathernext.py b/examples/weathernext.py
index 082b9081d..bc373f92c 100644
--- a/examples/weathernext.py
+++ b/examples/weathernext.py
@@ -126,9 +126,15 @@ def forward(
         # Encode each lat/lon cell independently, then flatten the grid to points:
         # [B, lat, lon, hidden] -> [B, grid_points, hidden].
         grid_latent = op.Tanh(self.grid_encoder(op, grid_features))
+        batch_dim = op.Shape(grid_latent, start=0, end=1)
+        flat_grid_shape = op.Concat(
+            batch_dim,
+            op.Constant(value_ints=[s.grid_points, s.hidden_size]),
+            axis=0,
+        )
         grid_points = op.Reshape(
             grid_latent,
-            op.Constant(value_ints=[0, s.grid_points, s.hidden_size]),
+            flat_grid_shape,
         )
 
         # Aggregate grid points onto the mesh with a fixed sparse-style projection:
@@ -146,9 +152,14 @@ def forward(
 
         # Return a one-step forecast grid: [B, lat, lon, output_variables].
         forecast_points = self.grid_decoder(op, grid_points)
+        forecast_shape = op.Concat(
+            batch_dim,
+            op.Constant(value_ints=[s.lat, s.lon, s.output_variables]),
+            axis=0,
+        )
         return op.Reshape(
             forecast_points,
-            op.Constant(value_ints=[0, s.lat, s.lon, s.output_variables]),
+            forecast_shape,
         )
 
 

From 9a7b5a4c5fc6739882c072fad8d9361a8b5330ae Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Mon, 10 Aug 2026 06:29:00 +0000
Subject: [PATCH 06/18] Document WeatherNext demo config shim

Signed-off-by: GitHub <noreply@github.com>

Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
---
 examples/weathernext.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/examples/weathernext.py b/examples/weathernext.py
index bc373f92c..151032f2b 100644
--- a/examples/weathernext.py
+++ b/examples/weathernext.py
@@ -235,6 +235,9 @@ def build_weathernext_demo_package(shape: WeatherNextDemoShape, dtype: ir.DataTy
 
     rng = np.random.default_rng(20260810)
     module = WeatherNextGridMeshBlock(rng, shape)
+    # build_from_module currently accepts BaseModelConfig subclasses; ArchitectureConfig
+    # validates a few LLM fields even though this WeatherNext task only reads dtype.
+    # Use the smallest valid inert values: no vocabulary, one dummy attention head.
     config = ArchitectureConfig(
         vocab_size=0,
         hidden_size=shape.hidden_size,

From 1a008e67dc78a5852c294028560d4fb027ad95fc Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Mon, 10 Aug 2026 06:30:08 +0000
Subject: [PATCH 07/18] Keep WeatherNext validation NumPy compatible

Signed-off-by: GitHub <noreply@github.com>

Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
---
 examples/weathernext.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/examples/weathernext.py b/examples/weathernext.py
index 151032f2b..bb9376b2d 100644
--- a/examples/weathernext.py
+++ b/examples/weathernext.py
@@ -271,14 +271,14 @@ def _validate_with_ort(output_dir: str, shape: WeatherNextDemoShape) -> None:
     rng = np.random.default_rng(42)
     feeds = {
         "input_state": rng.standard_normal(
-            (1, shape.lat, shape.lon, shape.input_variables), dtype=np.float32
-        ),
+            (1, shape.lat, shape.lon, shape.input_variables)
+        ).astype(np.float32),
         "forcings": rng.standard_normal(
-            (1, shape.lat, shape.lon, shape.forcing_variables), dtype=np.float32
-        ),
+            (1, shape.lat, shape.lon, shape.forcing_variables)
+        ).astype(np.float32),
         "sample_noise": rng.standard_normal(
-            (1, shape.lat, shape.lon, shape.noise_channels), dtype=np.float32
-        ),
+            (1, shape.lat, shape.lon, shape.noise_channels)
+        ).astype(np.float32),
     }
     (next_state,) = sess.run(None, feeds)
     print(f"Validation output next_state shape: {next_state.shape}")

From 8c98c0ec8e4f038cb1df62de4aad2717b178b77a Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Mon, 10 Aug 2026 06:31:16 +0000
Subject: [PATCH 08/18] Clarify WeatherNext demo graph comments

Signed-off-by: GitHub <noreply@github.com>

Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
---
 examples/weathernext.py | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/examples/weathernext.py b/examples/weathernext.py
index bb9376b2d..f1a08a96b 100644
--- a/examples/weathernext.py
+++ b/examples/weathernext.py
@@ -137,16 +137,18 @@ def forward(
             flat_grid_shape,
         )
 
-        # Aggregate grid points onto the mesh with a fixed sparse-style projection:
-        # [mesh_points, grid_points] @ [B, grid_points, hidden] -> [B, mesh_points, hidden].
+        # Aggregate grid points onto the mesh with a fixed sparse-style projection.
+        # MatMul broadcasts the 2-D projection over batch:
+        # [mesh_nodes, grid_points] @ [B, grid_points, hidden] -> [B, mesh_nodes, hidden].
         mesh_latent = op.MatMul(self.grid_to_mesh, grid_points)
 
         # A compact MLP stands in for WeatherNext's mesh GNN/update blocks.
         mesh_delta = op.Tanh(self.mesh_update_in(op, mesh_latent))
         mesh_latent = op.Add(mesh_latent, self.mesh_update_out(op, mesh_delta))
 
-        # Decode mesh latents back onto the lat/lon grid and add the encoded-grid residual:
-        # [grid_points, mesh_points] @ [B, mesh_points, hidden] -> [B, grid_points, hidden].
+        # Decode mesh latents back onto the lat/lon grid and add the encoded-grid residual.
+        # MatMul again broadcasts the 2-D projection over batch:
+        # [grid_points, mesh_nodes] @ [B, mesh_nodes, hidden] -> [B, grid_points, hidden].
         grid_delta = op.MatMul(self.mesh_to_grid, mesh_latent)
         grid_points = op.Add(grid_points, grid_delta)
 
@@ -235,9 +237,11 @@ def build_weathernext_demo_package(shape: WeatherNextDemoShape, dtype: ir.DataTy
 
     rng = np.random.default_rng(20260810)
     module = WeatherNextGridMeshBlock(rng, shape)
-    # build_from_module currently accepts BaseModelConfig subclasses; ArchitectureConfig
-    # validates a few LLM fields even though this WeatherNext task only reads dtype.
-    # Use the smallest valid inert values: no vocabulary, one dummy attention head.
+    # build_from_module currently accepts BaseModelConfig subclasses; this task reads only
+    # config.dtype, while build_from_module validates and uses dtype to cast parameters.
+    # The remaining ArchitectureConfig fields are inert validation shims. A production
+    # WeatherNext port should replace this with a dedicated task/config pair that records
+    # grid resolution, variable metadata, and mesh topology instead of LLM placeholders.
     config = ArchitectureConfig(
         vocab_size=0,
         hidden_size=shape.hidden_size,

From 2ce33ccd2261158b785118d1d5c80ab9ca7f0853 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Mon, 10 Aug 2026 13:57:16 +0000
Subject: [PATCH 09/18] Add formal WeatherNext model support

Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
---
 examples/weathernext.py                | 442 +++++++++----------------
 src/mobius/__init__.py                 |   8 +-
 src/mobius/_configs/__init__.py        |   2 +
 src/mobius/_configs/_weathernext.py    |  62 ++++
 src/mobius/integrations/weathernext.py | 192 +++++++++++
 src/mobius/models/__init__.py          |   2 +
 src/mobius/models/weathernext.py       | 137 ++++++++
 src/mobius/models/weathernext_test.py  | 110 ++++++
 src/mobius/tasks/__init__.py           |   3 +
 src/mobius/tasks/_weathernext.py       |  61 ++++
 tests/weathernext_example_test.py      |  51 +--
 11 files changed, 758 insertions(+), 312 deletions(-)
 create mode 100644 src/mobius/_configs/_weathernext.py
 create mode 100644 src/mobius/integrations/weathernext.py
 create mode 100644 src/mobius/models/weathernext.py
 create mode 100644 src/mobius/models/weathernext_test.py
 create mode 100644 src/mobius/tasks/_weathernext.py

diff --git a/examples/weathernext.py b/examples/weathernext.py
index f1a08a96b..e6af71f43 100644
--- a/examples/weathernext.py
+++ b/examples/weathernext.py
@@ -2,35 +2,29 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
-"""WeatherNext-style ONNX conversion demo.
+"""Build and run a WeatherNext-style one-step forecast ONNX model.
 
-WeatherNext 2 is a JAX/Haiku + xarray model rather than a HuggingFace
-``transformers`` model, so it does not fit the normal ``mobius build`` path.
-This example demonstrates the ONNX workflow for that family of models by
-defining the same high-level one-step forecast contract used by WeatherNext:
+The formal Mobius support lives in ``mobius.models.WeatherNextModel``,
+``mobius.tasks.WeatherNextForecastTask``, and ``mobius.integrations.weathernext``.
+This script demonstrates that path and can run the exported model on either:
 
-``input weather grid + forcings + stochastic noise → next weather grid``.
+* an ``.npz`` file with ``input_state``, ``forcings``, and ``sample_noise`` arrays, or
+* a local xarray NetCDF/Zarr weather dataset plus selected variable names.
 
-The tiny model below uses WeatherNext-like graph data flow:
-
-1. encode lat/lon grid variables at each grid cell,
-2. aggregate grid cells onto a mesh,
-3. update the mesh latent state,
-4. project mesh latents back to the grid, and
-5. decode per-grid-cell forecast variables.
-
-It intentionally uses small deterministic weights so the demo can be run
-without downloading WeatherNext checkpoints.  The graph I/O and component
-boundaries are the pieces to keep when replacing the toy modules with a full
-translation of google-deepmind/weathernext's Haiku modules and ``.npz``
-checkpoints.
+If no data or checkpoint is provided, the script uses deterministic demo weights
+and synthetic inputs so the ONNX workflow remains runnable in a fresh checkout.
 
 Usage::
 
-    python examples/weathernext.py output/weathernext-mini --validate
+    PYTHONPATH=src python examples/weathernext.py output/weathernext-mini --validate
+
+    PYTHONPATH=src python examples/weathernext.py output/weathernext-era5 \
+        --input-data era5_sample.npz --weights converted_weathernext_weights.npz --run
 
-    python examples/weathernext.py output/weathernext-mini \
-        --lat 8 --lon 16 --mesh-nodes 12 --hidden-size 32
+    PYTHONPATH=src python examples/weathernext.py output/weathernext-xarray \
+        --input-data weatherbench_sample.zarr \
+        --input-variable-names 2m_temperature mean_sea_level_pressure \
+        --forcing-variable-names toa_incident_solar_radiation --run
 """
 
 from __future__ import annotations
@@ -38,315 +32,185 @@
 import argparse
 import os
 import sys
-from dataclasses import dataclass
 
 import numpy as np
 import onnx_ir as ir
-from onnxscript import GraphBuilder, OpBuilder, nn
-
-import mobius
-from mobius import OPSET_VERSION, ArchitectureConfig, ModelPackage, build_from_module
-from mobius.tasks import ModelTask
-
-
-@dataclass(frozen=True)
-class WeatherNextDemoShape:
-    """Concrete shape for the one-step WeatherNext forecast demo."""
-
-    lat: int
-    lon: int
-    mesh_nodes: int
-    input_variables: int
-    forcing_variables: int
-    noise_channels: int
-    output_variables: int
-    hidden_size: int
-
-    @property
-    def grid_points(self) -> int:
-        return self.lat * self.lon
-
-    @property
-    def encoder_channels(self) -> int:
-        return self.input_variables + self.forcing_variables + self.noise_channels
-
-
-def _make_parameter(rng: np.random.Generator, shape: tuple[int, ...]) -> nn.Parameter:
-    values = rng.standard_normal(shape).astype(np.float32) * 0.05
-    return nn.Parameter(list(shape), data=ir.tensor(values))
-
 
-class DemoLinear(nn.Module):
-    """Small deterministic ``Linear`` layer for a self-contained runnable demo."""
+from mobius import WeatherNextConfig
+from mobius.integrations.weathernext import (
+    build_weathernext_package,
+    infer_config_from_feeds,
+    load_npz_forecast_inputs,
+    load_npz_weights,
+    load_xarray_forecast_inputs,
+)
 
-    def __init__(self, rng: np.random.Generator, in_features: int, out_features: int):
-        super().__init__()
-        self.weight = _make_parameter(rng, (out_features, in_features))
-        self.bias = _make_parameter(rng, (out_features,))
 
-    def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value:
-        # Project the trailing feature dimension: [..., in_features] -> [..., out_features].
-        x = op.MatMul(x, op.Transpose(self.weight, perm=[1, 0]))
-        return op.Add(x, self.bias)
-
-
-class WeatherNextGridMeshBlock(nn.Module):
-    """One grid→mesh→grid block mirroring WeatherNext's graph-forecast data flow."""
-
-    def __init__(self, rng: np.random.Generator, shape: WeatherNextDemoShape):
-        super().__init__()
-        self._shape = shape
-        self.grid_encoder = DemoLinear(rng, shape.encoder_channels, shape.hidden_size)
-        self.mesh_update_in = DemoLinear(rng, shape.hidden_size, 4 * shape.hidden_size)
-        self.mesh_update_out = DemoLinear(rng, 4 * shape.hidden_size, shape.hidden_size)
-        self.grid_decoder = DemoLinear(rng, shape.hidden_size, shape.output_variables)
-
-        grid_to_mesh = _projection_matrix(shape.mesh_nodes, shape.grid_points)
-        mesh_to_grid = _projection_matrix(shape.grid_points, shape.mesh_nodes)
-        self.grid_to_mesh = nn.Parameter(
-            [shape.mesh_nodes, shape.grid_points], data=ir.tensor(grid_to_mesh)
-        )
-        self.mesh_to_grid = nn.Parameter(
-            [shape.grid_points, shape.mesh_nodes], data=ir.tensor(mesh_to_grid)
-        )
+def _resolve_dtype(name: str) -> ir.DataType:
+    if name == "f32":
+        return ir.DataType.FLOAT
+    if name == "f16":
+        return ir.DataType.FLOAT16
+    raise ValueError(f"Unsupported dtype: {name}")
 
-    def forward(
-        self,
-        op: OpBuilder,
-        input_state: ir.Value,
-        forcings: ir.Value,
-        sample_noise: ir.Value,
-    ) -> ir.Value:
-        s = self._shape
-
-        # Concatenate per-cell weather variables, known future forcings, and FGN noise:
-        # [B, lat, lon, input+forcing+noise].
-        grid_features = op.Concat(input_state, forcings, sample_noise, axis=-1)
-
-        # Encode each lat/lon cell independently, then flatten the grid to points:
-        # [B, lat, lon, hidden] -> [B, grid_points, hidden].
-        grid_latent = op.Tanh(self.grid_encoder(op, grid_features))
-        batch_dim = op.Shape(grid_latent, start=0, end=1)
-        flat_grid_shape = op.Concat(
-            batch_dim,
-            op.Constant(value_ints=[s.grid_points, s.hidden_size]),
-            axis=0,
-        )
-        grid_points = op.Reshape(
-            grid_latent,
-            flat_grid_shape,
-        )
 
-        # Aggregate grid points onto the mesh with a fixed sparse-style projection.
-        # MatMul broadcasts the 2-D projection over batch:
-        # [mesh_nodes, grid_points] @ [B, grid_points, hidden] -> [B, mesh_nodes, hidden].
-        mesh_latent = op.MatMul(self.grid_to_mesh, grid_points)
-
-        # A compact MLP stands in for WeatherNext's mesh GNN/update blocks.
-        mesh_delta = op.Tanh(self.mesh_update_in(op, mesh_latent))
-        mesh_latent = op.Add(mesh_latent, self.mesh_update_out(op, mesh_delta))
-
-        # Decode mesh latents back onto the lat/lon grid and add the encoded-grid residual.
-        # MatMul again broadcasts the 2-D projection over batch:
-        # [grid_points, mesh_nodes] @ [B, mesh_nodes, hidden] -> [B, grid_points, hidden].
-        grid_delta = op.MatMul(self.mesh_to_grid, mesh_latent)
-        grid_points = op.Add(grid_points, grid_delta)
-
-        # Return a one-step forecast grid: [B, lat, lon, output_variables].
-        forecast_points = self.grid_decoder(op, grid_points)
-        forecast_shape = op.Concat(
-            batch_dim,
-            op.Constant(value_ints=[s.lat, s.lon, s.output_variables]),
-            axis=0,
+def _load_real_data(args: argparse.Namespace) -> dict[str, np.ndarray] | None:
+    if args.input_data is None:
+        return None
+    if args.input_data.endswith(".npz"):
+        return load_npz_forecast_inputs(args.input_data)
+    if not args.input_variable_names or not args.forcing_variable_names:
+        raise ValueError(
+            "--input-variable-names and --forcing-variable-names are required for xarray data"
         )
-        return op.Reshape(
-            forecast_points,
-            forecast_shape,
-        )
-
-
-class WeatherNextDemoTask(ModelTask):
-    """Task wiring for a one-step WeatherNext-style forecast graph."""
-
-    model_roles = {"model": "forecast"}
+    return load_xarray_forecast_inputs(
+        args.input_data,
+        input_variables=args.input_variable_names,
+        forcing_variables=args.forcing_variable_names,
+        noise_channels=args.noise_channels,
+        batch_index=args.batch_index,
+        sample_noise_seed=args.sample_noise_seed,
+    )
 
-    def __init__(self, shape: WeatherNextDemoShape):
-        self._shape = shape
 
-    def build(self, module: nn.Module, config: ArchitectureConfig) -> ModelPackage:
-        batch = ir.SymbolicDim("batch")
-        s = self._shape
+def _synthetic_feeds(config: WeatherNextConfig) -> dict[str, np.ndarray]:
+    rng = np.random.default_rng(42)
+    return {
+        "input_state": rng.standard_normal(
+            (1, config.lat, config.lon, config.input_variables)
+        ).astype(np.float32),
+        "forcings": rng.standard_normal(
+            (1, config.lat, config.lon, config.forcing_variables)
+        ).astype(np.float32),
+        "sample_noise": rng.standard_normal(
+            (1, config.lat, config.lon, config.noise_channels)
+        ).astype(np.float32),
+    }
 
-        graph, builder = _make_demo_graph("weathernext_one_step_forecast")
-        op = builder.op
 
-        input_state = builder.input(
-            "input_state",
-            dtype=config.dtype,
-            shape=[batch, s.lat, s.lon, s.input_variables],
-        )
-        forcings = builder.input(
-            "forcings",
-            dtype=config.dtype,
-            shape=[batch, s.lat, s.lon, s.forcing_variables],
+def _config_from_args(args: argparse.Namespace, feeds: dict[str, np.ndarray] | None):
+    dtype = _resolve_dtype(args.dtype)
+    if feeds is not None:
+        return infer_config_from_feeds(
+            feeds,
+            mesh_nodes=args.mesh_nodes,
+            hidden_size=args.hidden_size,
+            intermediate_size=args.intermediate_size,
+            num_hidden_layers=args.num_hidden_layers,
+            dtype=dtype,
         )
-        sample_noise = builder.input(
-            "sample_noise",
-            dtype=config.dtype,
-            shape=[batch, s.lat, s.lon, s.noise_channels],
-        )
-
-        next_state = module(op, input_state, forcings, sample_noise)
-        builder.add_output(next_state, "next_state")
-
-        return ModelPackage({"model": _make_demo_model(graph)}, config=config)
-
-
-def _make_demo_graph(name: str) -> tuple[ir.Graph, GraphBuilder]:
-    graph = ir.Graph(
-        [],
-        [],
-        nodes=[],
-        name=name,
-        opset_imports={"": OPSET_VERSION, "com.microsoft": 1},
-    )
-    return graph, GraphBuilder(graph)
-
-
-def _make_demo_model(graph: ir.Graph) -> ir.Model:
-    model = ir.Model(graph, ir_version=11)
-    model.producer_name = "mobius"
-    model.producer_version = mobius.__version__
-    return model
-
-
-def _projection_matrix(rows: int, cols: int) -> np.ndarray:
-    """Create deterministic normalized projections between grid and mesh points."""
-
-    row_positions = np.linspace(0.0, 1.0, rows, dtype=np.float32)[:, None]
-    col_positions = np.linspace(0.0, 1.0, cols, dtype=np.float32)[None, :]
-    distance = np.abs(row_positions - col_positions)
-    weights = np.maximum(1.0 - 2.0 * distance, 0.0)
-    weights += 1e-3
-    weights /= weights.sum(axis=1, keepdims=True)
-    return weights.astype(np.float32)
-
-
-def build_weathernext_demo_package(shape: WeatherNextDemoShape, dtype: ir.DataType) -> ModelPackage:
-    """Build the demo WeatherNext-style ONNX package."""
-
-    rng = np.random.default_rng(20260810)
-    module = WeatherNextGridMeshBlock(rng, shape)
-    # build_from_module currently accepts BaseModelConfig subclasses; this task reads only
-    # config.dtype, while build_from_module validates and uses dtype to cast parameters.
-    # The remaining ArchitectureConfig fields are inert validation shims. A production
-    # WeatherNext port should replace this with a dedicated task/config pair that records
-    # grid resolution, variable metadata, and mesh topology instead of LLM placeholders.
-    config = ArchitectureConfig(
-        vocab_size=0,
-        hidden_size=shape.hidden_size,
-        intermediate_size=4 * shape.hidden_size,
-        num_hidden_layers=1,
-        num_attention_heads=1,
-        num_key_value_heads=1,
-        head_dim=shape.hidden_size,
+    return WeatherNextConfig(
+        lat=args.lat,
+        lon=args.lon,
+        mesh_nodes=args.mesh_nodes,
+        input_variables=args.input_variables,
+        forcing_variables=args.forcing_variables,
+        noise_channels=args.noise_channels,
+        output_variables=args.output_variables,
+        hidden_size=args.hidden_size,
+        intermediate_size=args.intermediate_size or 4 * args.hidden_size,
+        num_hidden_layers=args.num_hidden_layers,
         dtype=dtype,
     )
-    return build_from_module(module, config, task=WeatherNextDemoTask(shape))
 
 
-def _resolve_dtype(name: str) -> ir.DataType:
-    if name == "f32":
-        return ir.DataType.FLOAT
-    if name == "f16":
-        return ir.DataType.FLOAT16
-    raise ValueError(f"Unsupported dtype: {name}")
-
-
-def _validate_with_ort(output_dir: str, shape: WeatherNextDemoShape) -> None:
+def _run_with_ort(output_dir: str, feeds: dict[str, np.ndarray]) -> None:
     try:
         import onnxruntime as ort
     except ImportError:
-        print("onnxruntime is not installed; skipping validation.", file=sys.stderr)
+        print("onnxruntime is not installed; skipping inference.", file=sys.stderr)
         return
 
     model_path = os.path.join(output_dir, "model.onnx")
     sess = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
-    rng = np.random.default_rng(42)
-    feeds = {
-        "input_state": rng.standard_normal(
-            (1, shape.lat, shape.lon, shape.input_variables)
-        ).astype(np.float32),
-        "forcings": rng.standard_normal(
-            (1, shape.lat, shape.lon, shape.forcing_variables)
-        ).astype(np.float32),
-        "sample_noise": rng.standard_normal(
-            (1, shape.lat, shape.lon, shape.noise_channels)
-        ).astype(np.float32),
-    }
-    (next_state,) = sess.run(None, feeds)
-    print(f"Validation output next_state shape: {next_state.shape}")
-    print(f"Validation output range: [{next_state.min():.6f}, {next_state.max():.6f}]")
+    (next_state,) = sess.run(["next_state"], feeds)
+    print(f"Inference output next_state shape: {next_state.shape}")
+    print(f"Inference output range: [{next_state.min():.6f}, {next_state.max():.6f}]")
 
 
 def _parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(
-        description="Build a runnable WeatherNext-style grid→mesh→grid ONNX demo.",
+        description="Build a Mobius WeatherNext one-step forecast ONNX graph.",
     )
     parser.add_argument("output_dir", help="Directory to save model.onnx and model.onnx.data.")
-    parser.add_argument("--lat", type=int, default=4, help="Number of latitude points.")
-    parser.add_argument("--lon", type=int, default=8, help="Number of longitude points.")
-    parser.add_argument("--mesh-nodes", type=int, default=6, help="Number of demo mesh nodes.")
-    parser.add_argument("--input-variables", type=int, default=5, help="Input weather channels.")
-    parser.add_argument("--forcing-variables", type=int, default=2, help="Known forcing channels.")
-    parser.add_argument("--noise-channels", type=int, default=2, help="FGN stochastic noise channels.")
-    parser.add_argument("--output-variables", type=int, default=5, help="Forecast weather channels.")
+    parser.add_argument("--lat", type=int, default=4, help="Synthetic-data latitude points.")
+    parser.add_argument("--lon", type=int, default=8, help="Synthetic-data longitude points.")
+    parser.add_argument("--mesh-nodes", type=int, default=6, help="Forecast mesh nodes.")
+    parser.add_argument(
+        "--input-variables", type=int, default=5, help="Synthetic input channels."
+    )
+    parser.add_argument(
+        "--forcing-variables", type=int, default=2, help="Synthetic forcing channels."
+    )
+    parser.add_argument(
+        "--noise-channels", type=int, default=2, help="Stochastic noise channels."
+    )
+    parser.add_argument(
+        "--output-variables", type=int, default=5, help="Synthetic output channels."
+    )
     parser.add_argument("--hidden-size", type=int, default=16, help="Latent feature size.")
-    parser.add_argument("--dtype", choices=["f32", "f16"], default="f32", help="ONNX weight dtype.")
-    parser.add_argument("--validate", action="store_true", help="Run one ONNX Runtime inference.")
+    parser.add_argument("--intermediate-size", type=int, help="Mesh MLP intermediate size.")
+    parser.add_argument("--num-hidden-layers", type=int, default=1, help="Mesh update blocks.")
+    parser.add_argument(
+        "--dtype", choices=["f32", "f16"], default="f32", help="ONNX weight dtype."
+    )
+    parser.add_argument("--weights", help="Optional Mobius-aligned WeatherNext weights .npz.")
+    parser.add_argument(
+        "--input-data",
+        help="Optional .npz, NetCDF, or Zarr weather sample used for real-data inference.",
+    )
+    parser.add_argument(
+        "--input-variable-names",
+        nargs="+",
+        help="xarray variables stacked into input_state channels.",
+    )
+    parser.add_argument(
+        "--forcing-variable-names",
+        nargs="+",
+        help="xarray variables stacked into forcings channels.",
+    )
+    parser.add_argument(
+        "--batch-index", type=int, default=0, help="Time/batch index for xarray."
+    )
+    parser.add_argument(
+        "--sample-noise-seed", type=int, default=0, help="Generated noise seed."
+    )
+    parser.add_argument("--run", action="store_true", help="Run one ONNX Runtime inference.")
+    parser.add_argument(
+        "--validate",
+        action="store_true",
+        help="Alias for --run, kept for the original self-contained demo workflow.",
+    )
     return parser.parse_args()
 
 
 def main() -> None:
     args = _parse_args()
-    shape = WeatherNextDemoShape(
-        lat=args.lat,
-        lon=args.lon,
-        mesh_nodes=args.mesh_nodes,
-        input_variables=args.input_variables,
-        forcing_variables=args.forcing_variables,
-        noise_channels=args.noise_channels,
-        output_variables=args.output_variables,
-        hidden_size=args.hidden_size,
+    feeds = _load_real_data(args)
+    config = _config_from_args(args, feeds)
+    config.validate()
+
+    print("Building WeatherNext one-step forecast ONNX graph...")
+    print(f"  grid: {config.lat} x {config.lon} ({config.grid_points} cells)")
+    print(f"  mesh nodes: {config.mesh_nodes}")
+    print(
+        "  channels: "
+        f"input={config.input_variables}, forcing={config.forcing_variables}, "
+        f"noise={config.noise_channels}, output={config.output_variables}"
     )
 
-    if min(
-        shape.lat,
-        shape.lon,
-        shape.mesh_nodes,
-        shape.input_variables,
-        shape.forcing_variables,
-        shape.noise_channels,
-        shape.output_variables,
-        shape.hidden_size,
-    ) <= 0:
-        raise ValueError("All shape arguments must be positive.")
-
-    print("Building WeatherNext-style one-step forecast ONNX graph...")
-    print(f"  grid: {shape.lat} x {shape.lon} ({shape.grid_points} cells)")
-    print(f"  mesh nodes: {shape.mesh_nodes}")
-    print(f"  channels: input={shape.input_variables}, forcing={shape.forcing_variables}, "
-          f"noise={shape.noise_channels}, output={shape.output_variables}")
-
-    pkg = build_weathernext_demo_package(shape, dtype=_resolve_dtype(args.dtype))
-    model = pkg["model"]
+    weights = load_npz_weights(args.weights) if args.weights else None
+    package = build_weathernext_package(config, weights=weights)
+    model = package["model"]
     print(f"Built model with {model.graph.num_nodes()} ONNX nodes.")
 
-    pkg.save(args.output_dir, check_weights=True, progress_bar=False)
-    print(f"Saved WeatherNext demo package to {args.output_dir!r}.")
+    package.save(args.output_dir, check_weights=True, progress_bar=False)
+    print(f"Saved WeatherNext package to {args.output_dir!r}.")
 
-    if args.validate:
-        _validate_with_ort(args.output_dir, shape)
+    if args.run or args.validate:
+        _run_with_ort(
+            args.output_dir, feeds if feeds is not None else _synthetic_feeds(config)
+        )
 
 
 if __name__ == "__main__":
diff --git a/src/mobius/__init__.py b/src/mobius/__init__.py
index 0962f505a..373e24020 100644
--- a/src/mobius/__init__.py
+++ b/src/mobius/__init__.py
@@ -30,6 +30,9 @@
     "VisionConfig",
     "VisionLanguageConfig",
     "WhisperConfig",
+    "WeatherNextConfig",
+    "WeatherNextForecastTask",
+    "WeatherNextModel",
     "WorldModelConfig",
     "WorldModelTask",
     "YolosConfig",
@@ -78,6 +81,7 @@
     SegformerConfig,
     VisionConfig,
     VisionLanguageConfig,
+    WeatherNextConfig,
     WhisperConfig,
     WorldModelConfig,
     YolosConfig,
@@ -95,5 +99,5 @@
 from mobius._weight_loading import apply_weights
 from mobius.integrations.gguf import build_from_gguf
 from mobius.integrations.nemo import build_from_nemo
-from mobius.models import MLPWorldModel
-from mobius.tasks import CausalLMTask, ModelTask, WorldModelTask
+from mobius.models import MLPWorldModel, WeatherNextModel
+from mobius.tasks import CausalLMTask, ModelTask, WeatherNextForecastTask, WorldModelTask
diff --git a/src/mobius/_configs/__init__.py b/src/mobius/_configs/__init__.py
index afb241a16..d5e30e945 100644
--- a/src/mobius/_configs/__init__.py
+++ b/src/mobius/_configs/__init__.py
@@ -79,6 +79,7 @@
     TTSConfig,
     VisionConfig,
 )
+from mobius._configs._weathernext import WeatherNextConfig
 from mobius._configs._world_model import WorldModelConfig
 
 __all__ = [
@@ -118,6 +119,7 @@
     "VisionConfig",
     "VisionLanguageConfig",
     "WhisperConfig",
+    "WeatherNextConfig",
     "WorldModelConfig",
     "YolosConfig",
     "Zamba2Config",
diff --git a/src/mobius/_configs/_weathernext.py b/src/mobius/_configs/_weathernext.py
new file mode 100644
index 000000000..f38db7149
--- /dev/null
+++ b/src/mobius/_configs/_weathernext.py
@@ -0,0 +1,62 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Configuration for WeatherNext-style one-step forecast graphs."""
+
+from __future__ import annotations
+
+import dataclasses
+
+from mobius._configs._base import BaseModelConfig
+
+
+@dataclasses.dataclass
+class WeatherNextConfig(BaseModelConfig):
+    """Configuration for grid→mesh→grid WeatherNext forecast modules.
+
+    Shapes exclude the leading batch dimension. The task exposes a one-step
+    forecast contract:
+    ``input_state + forcings + sample_noise -> next_state``.
+    """
+
+    lat: int = 4
+    lon: int = 8
+    mesh_nodes: int = 6
+    input_variables: int = 5
+    forcing_variables: int = 2
+    noise_channels: int = 2
+    output_variables: int = 5
+    hidden_size: int = 16
+    intermediate_size: int = 64
+    num_hidden_layers: int = 1
+    hidden_act: str | None = "silu"
+
+    @property
+    def grid_points(self) -> int:
+        """Number of lat/lon grid cells."""
+        return self.lat * self.lon
+
+    @property
+    def encoder_channels(self) -> int:
+        """Per-grid-cell input channels after concatenating all inputs."""
+        return self.input_variables + self.forcing_variables + self.noise_channels
+
+    def validate(self) -> None:
+        """Validate dimensions required by the WeatherNext forecast task."""
+        for name in (
+            "lat",
+            "lon",
+            "mesh_nodes",
+            "input_variables",
+            "forcing_variables",
+            "noise_channels",
+            "output_variables",
+            "hidden_size",
+            "intermediate_size",
+            "num_hidden_layers",
+        ):
+            value = getattr(self, name)
+            if not isinstance(value, int) or isinstance(value, bool) or value <= 0:
+                raise ValueError(f"{name} must be a positive integer")
+        if self.hidden_act is None:
+            raise ValueError("hidden_act must be set")
diff --git a/src/mobius/integrations/weathernext.py b/src/mobius/integrations/weathernext.py
new file mode 100644
index 000000000..78fa1c096
--- /dev/null
+++ b/src/mobius/integrations/weathernext.py
@@ -0,0 +1,192 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Helpers for building and running WeatherNext-style Mobius packages."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import onnx_ir as ir
+import torch
+
+from mobius._builder import build_from_module
+from mobius._configs import WeatherNextConfig
+from mobius._model_package import ModelPackage
+from mobius.models import WeatherNextModel
+
+
+def load_npz_weights(path: str | Path) -> dict[str, torch.Tensor]:
+    """Load a Mobius-aligned WeatherNext state dict from an ``.npz`` file."""
+    with np.load(path) as data:
+        return {name: torch.from_numpy(np.asarray(data[name])) for name in data.files}
+
+
+def create_demo_state_dict(
+    config: WeatherNextConfig,
+    *,
+    seed: int = 20260810,
+) -> dict[str, torch.Tensor]:
+    """Create deterministic weights for examples and tests.
+
+    These weights are intentionally small and are not a trained WeatherNext
+    checkpoint. Pass ``weights=load_npz_weights(...)`` to
+    :func:`build_weathernext_package` for converted checkpoint weights.
+    """
+    rng = np.random.default_rng(seed)
+
+    def parameter(shape: tuple[int, ...]) -> torch.Tensor:
+        return torch.from_numpy(rng.standard_normal(shape).astype(np.float32) * 0.05)
+
+    state: dict[str, torch.Tensor] = {
+        "grid_encoder.weight": parameter((config.hidden_size, config.encoder_channels)),
+        "grid_encoder.bias": parameter((config.hidden_size,)),
+        "grid_decoder.weight": parameter((config.output_variables, config.hidden_size)),
+        "grid_decoder.bias": parameter((config.output_variables,)),
+    }
+    for layer_idx in range(config.num_hidden_layers):
+        state[f"mesh_update_in.{layer_idx}.weight"] = parameter(
+            (config.intermediate_size, config.hidden_size)
+        )
+        state[f"mesh_update_in.{layer_idx}.bias"] = parameter((config.intermediate_size,))
+        state[f"mesh_update_out.{layer_idx}.weight"] = parameter(
+            (config.hidden_size, config.intermediate_size)
+        )
+        state[f"mesh_update_out.{layer_idx}.bias"] = parameter((config.hidden_size,))
+    return state
+
+
+def build_weathernext_package(
+    config: WeatherNextConfig,
+    *,
+    weights: dict[str, torch.Tensor] | None = None,
+    execution_provider: str = "default",
+) -> ModelPackage:
+    """Build a WeatherNext one-step forecast package and apply weights."""
+    package = build_from_module(
+        WeatherNextModel(config),
+        config,
+        task="weathernext-forecast",
+        execution_provider=execution_provider,
+    )
+    package.apply_weights(weights if weights is not None else create_demo_state_dict(config))
+    return package
+
+
+def load_npz_forecast_inputs(path: str | Path) -> dict[str, np.ndarray]:
+    """Load ``input_state``, ``forcings``, and ``sample_noise`` arrays from ``.npz``."""
+    with np.load(path) as data:
+        feeds = {name: np.asarray(data[name], dtype=np.float32) for name in _INPUT_NAMES}
+    return feeds
+
+
+def load_xarray_forecast_inputs(
+    path: str | Path,
+    *,
+    input_variables: list[str],
+    forcing_variables: list[str],
+    noise_channels: int,
+    batch_index: int = 0,
+    sample_noise_seed: int = 0,
+) -> dict[str, np.ndarray]:
+    """Load WeatherNext inputs from a local xarray NetCDF or Zarr dataset.
+
+    The selected variables must share latitude/longitude dimensions. Optional
+    time dimensions are indexed by ``batch_index`` and each variable is stacked
+    into the trailing channel dimension expected by the ONNX graph.
+    """
+    try:
+        import xarray as xr
+    except ImportError as e:  # pragma: no cover - exercised only without optional xarray
+        raise ImportError("xarray is required to read NetCDF/Zarr WeatherNext inputs") from e
+
+    path = Path(path)
+    dataset = (
+        xr.open_zarr(path)
+        if path.is_dir() or path.suffix == ".zarr"
+        else xr.open_dataset(path)
+    )
+    try:
+        input_state = _stack_xarray_variables(
+            dataset,
+            input_variables,
+            batch_index=batch_index,
+        )
+        forcings = _stack_xarray_variables(
+            dataset,
+            forcing_variables,
+            batch_index=batch_index,
+        )
+    finally:
+        dataset.close()
+
+    rng = np.random.default_rng(sample_noise_seed)
+    sample_noise = rng.standard_normal(
+        (input_state.shape[0], input_state.shape[1], input_state.shape[2], noise_channels)
+    ).astype(np.float32)
+    return {
+        "input_state": input_state,
+        "forcings": forcings,
+        "sample_noise": sample_noise,
+    }
+
+
+def infer_config_from_feeds(
+    feeds: dict[str, np.ndarray],
+    *,
+    mesh_nodes: int,
+    hidden_size: int,
+    intermediate_size: int | None = None,
+    num_hidden_layers: int = 1,
+    dtype: ir.DataType = ir.DataType.FLOAT,
+) -> WeatherNextConfig:
+    """Infer a WeatherNext config from loaded forecast input arrays."""
+    input_state = feeds["input_state"]
+    forcings = feeds["forcings"]
+    sample_noise = feeds["sample_noise"]
+    if input_state.ndim != 4 or forcings.ndim != 4 or sample_noise.ndim != 4:
+        raise ValueError(
+            "WeatherNext inputs must be rank-4 [batch, lat, lon, channels] arrays"
+        )
+    if (
+        input_state.shape[:3] != forcings.shape[:3]
+        or input_state.shape[:3] != sample_noise.shape[:3]
+    ):
+        raise ValueError("WeatherNext inputs must have matching batch/lat/lon dimensions")
+    return WeatherNextConfig(
+        lat=int(input_state.shape[1]),
+        lon=int(input_state.shape[2]),
+        mesh_nodes=mesh_nodes,
+        input_variables=int(input_state.shape[3]),
+        forcing_variables=int(forcings.shape[3]),
+        noise_channels=int(sample_noise.shape[3]),
+        output_variables=int(input_state.shape[3]),
+        hidden_size=hidden_size,
+        intermediate_size=intermediate_size or 4 * hidden_size,
+        num_hidden_layers=num_hidden_layers,
+        dtype=dtype,
+    )
+
+
+_INPUT_NAMES = ("input_state", "forcings", "sample_noise")
+
+
+def _stack_xarray_variables(
+    dataset: Any,
+    names: list[str],
+    *,
+    batch_index: int,
+) -> np.ndarray:
+    arrays = []
+    for name in names:
+        if name not in dataset:
+            raise KeyError(f"Variable {name!r} not found in dataset")
+        value = dataset[name]
+        for dim in value.dims:
+            if dim.lower() in {"time", "batch"}:
+                value = value.isel({dim: batch_index})
+        arrays.append(np.asarray(value, dtype=np.float32))
+    stacked = np.stack(arrays, axis=-1)
+    return stacked[None, ...] if stacked.ndim == 3 else stacked
diff --git a/src/mobius/models/__init__.py b/src/mobius/models/__init__.py
index 9ca9e2ce3..a1f6bb934 100644
--- a/src/mobius/models/__init__.py
+++ b/src/mobius/models/__init__.py
@@ -149,6 +149,7 @@
     "Wav2Vec2ForCTCModel",
     "Wav2Vec2Model",
     "WhisperForConditionalGeneration",
+    "WeatherNextModel",
     "MLPWorldModel",
     "XLMCausalLMModel",
     "Zamba2CausalLMModel",
@@ -309,6 +310,7 @@
 from mobius.models.vit import ViTModel
 from mobius.models.wav2vec2 import Wav2Vec2Model
 from mobius.models.wav2vec2_ctc import Wav2Vec2ForCTCModel
+from mobius.models.weathernext import WeatherNextModel
 from mobius.models.whisper import WhisperForConditionalGeneration
 from mobius.models.world_model import MLPWorldModel
 from mobius.models.xlm import XLMCausalLMModel
diff --git a/src/mobius/models/weathernext.py b/src/mobius/models/weathernext.py
new file mode 100644
index 000000000..bf945b7de
--- /dev/null
+++ b/src/mobius/models/weathernext.py
@@ -0,0 +1,137 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""WeatherNext-style grid→mesh→grid forecast model.
+
+This module provides Mobius-native ONNX graph construction for the one-step
+forecast contract used by WeatherNext-family models. It does not trace JAX;
+instead, each graph stage is declared directly with ONNX ops.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import onnx_ir as ir
+import torch
+from onnxscript import OpBuilder, nn
+
+from mobius._configs import WeatherNextConfig
+from mobius.components import Linear, get_activation
+
+
+def projection_matrix(rows: int, cols: int) -> np.ndarray:
+    """Create a normalized deterministic projection between grid and mesh points."""
+    row_positions = np.linspace(0.0, 1.0, rows, dtype=np.float32)[:, None]
+    col_positions = np.linspace(0.0, 1.0, cols, dtype=np.float32)[None, :]
+    distance = np.abs(row_positions - col_positions)
+    weights = np.maximum(1.0 - 2.0 * distance, 0.0)
+    weights += 1e-3
+    weights /= weights.sum(axis=1, keepdims=True)
+    return weights.astype(np.float32)
+
+
+class WeatherNextModel(nn.Module):
+    """One-step WeatherNext-style grid→mesh→grid forecast module.
+
+    Inputs:
+        - input_state: ``[batch, lat, lon, input_variables]``
+        - forcings: ``[batch, lat, lon, forcing_variables]``
+        - sample_noise: ``[batch, lat, lon, noise_channels]``
+
+    Output:
+        - next_state: ``[batch, lat, lon, output_variables]``
+    """
+
+    default_task = "weathernext-forecast"
+    config_class = WeatherNextConfig
+    category = "Weather"
+
+    def __init__(self, config: WeatherNextConfig):
+        super().__init__()
+        config.validate()
+        self.config = config
+        self.grid_encoder = Linear(config.encoder_channels, config.hidden_size)
+        self.mesh_update_in = nn.ModuleList(
+            [
+                Linear(config.hidden_size, config.intermediate_size)
+                for _ in range(config.num_hidden_layers)
+            ]
+        )
+        self.mesh_update_out = nn.ModuleList(
+            [
+                Linear(config.intermediate_size, config.hidden_size)
+                for _ in range(config.num_hidden_layers)
+            ]
+        )
+        self.grid_decoder = Linear(config.hidden_size, config.output_variables)
+        self._activation = get_activation(config.hidden_act)
+
+        # Fixed topology projections model the grid↔mesh connectivity. A real
+        # checkpoint may override these names with learned/sparse projection data.
+        self.grid_to_mesh = nn.Parameter(
+            [config.mesh_nodes, config.grid_points],
+            data=ir.tensor(projection_matrix(config.mesh_nodes, config.grid_points)),
+        )
+        self.mesh_to_grid = nn.Parameter(
+            [config.grid_points, config.mesh_nodes],
+            data=ir.tensor(projection_matrix(config.grid_points, config.mesh_nodes)),
+        )
+
+    def forward(
+        self,
+        op: OpBuilder,
+        input_state: ir.Value,
+        forcings: ir.Value,
+        sample_noise: ir.Value,
+    ) -> ir.Value:
+        config = self.config
+
+        # Concatenate per-cell weather variables, future forcings, and stochastic
+        # noise: [B, lat, lon, input+forcing+noise].
+        grid_features = op.Concat(input_state, forcings, sample_noise, axis=-1)
+
+        # Encode each lat/lon cell independently, then flatten the spatial grid:
+        # [B, lat, lon, hidden] -> [B, grid_points, hidden].
+        grid_latent = self._activation(op, self.grid_encoder(op, grid_features))
+        batch_dim = op.Shape(grid_latent, start=0, end=1)
+        flat_grid_shape = op.Concat(
+            batch_dim,
+            op.Constant(value_ints=[config.grid_points, config.hidden_size]),
+            axis=0,
+        )
+        grid_points = op.Reshape(grid_latent, flat_grid_shape)
+
+        # Aggregate encoded grid cells onto mesh nodes:
+        # [mesh_nodes, grid_points] @ [B, grid_points, hidden]
+        # -> [B, mesh_nodes, hidden].
+        mesh_latent = op.MatMul(self.grid_to_mesh, grid_points)
+
+        # Apply one or more residual mesh-update MLP blocks.
+        for update_in, update_out in zip(
+            self.mesh_update_in, self.mesh_update_out, strict=True
+        ):
+            mesh_delta = self._activation(op, update_in(op, mesh_latent))
+            mesh_latent = op.Add(mesh_latent, update_out(op, mesh_delta))
+
+        # Decode mesh latents back to grid cells and retain the encoded-grid
+        # residual: [grid_points, mesh_nodes] @ [B, mesh_nodes, hidden]
+        # -> [B, grid_points, hidden].
+        grid_delta = op.MatMul(self.mesh_to_grid, mesh_latent)
+        grid_points = op.Add(grid_points, grid_delta)
+
+        # Return one forecast step on the original grid:
+        # [B, grid_points, output_variables] -> [B, lat, lon, output_variables].
+        forecast_points = self.grid_decoder(op, grid_points)
+        forecast_shape = op.Concat(
+            batch_dim,
+            op.Constant(value_ints=[config.lat, config.lon, config.output_variables]),
+            axis=0,
+        )
+        return op.Reshape(forecast_points, forecast_shape)
+
+    def preprocess_weights(
+        self,
+        state_dict: dict[str, torch.Tensor],
+    ) -> dict[str, torch.Tensor]:
+        """Return WeatherNext weights unchanged after external integration mapping."""
+        return state_dict
diff --git a/src/mobius/models/weathernext_test.py b/src/mobius/models/weathernext_test.py
new file mode 100644
index 000000000..908e93cff
--- /dev/null
+++ b/src/mobius/models/weathernext_test.py
@@ -0,0 +1,110 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from __future__ import annotations
+
+import numpy as np
+import onnx_ir as ir
+import pytest
+
+from mobius import WeatherNextConfig, WeatherNextForecastTask, WeatherNextModel, build_from_module
+from mobius.integrations.weathernext import (
+    build_weathernext_package,
+    create_demo_state_dict,
+    infer_config_from_feeds,
+    load_npz_forecast_inputs,
+)
+from mobius.tasks import TASK_REGISTRY, get_task
+
+
+def _config() -> WeatherNextConfig:
+    return WeatherNextConfig(
+        lat=3,
+        lon=4,
+        mesh_nodes=5,
+        input_variables=2,
+        forcing_variables=1,
+        noise_channels=1,
+        output_variables=2,
+        hidden_size=8,
+        intermediate_size=16,
+        num_hidden_layers=2,
+    )
+
+
+class TestWeatherNextConfig:
+    def test_derived_sizes(self):
+        config = _config()
+        assert config.grid_points == 12
+        assert config.encoder_channels == 4
+
+    @pytest.mark.parametrize(
+        ("field", "value"),
+        [
+            ("lat", 0),
+            ("lon", True),
+            ("mesh_nodes", -1),
+            ("input_variables", 0),
+            ("forcing_variables", 0),
+            ("noise_channels", 0),
+            ("output_variables", 0),
+            ("hidden_size", 0),
+            ("intermediate_size", 0),
+            ("num_hidden_layers", 0),
+            ("hidden_act", None),
+        ],
+    )
+    def test_invalid_config_raises(self, field, value):
+        config = _config()
+        setattr(config, field, value)
+        with pytest.raises(ValueError):
+            config.validate()
+
+
+class TestWeatherNextForecastTask:
+    def test_registered(self):
+        assert TASK_REGISTRY["weathernext-forecast"] is WeatherNextForecastTask
+        assert isinstance(get_task("weathernext-forecast"), WeatherNextForecastTask)
+
+    def test_graph_contract(self, tmp_path):
+        config = _config()
+        package = build_weathernext_package(config)
+        package.save(tmp_path, check_weights=True, progress_bar=False)
+
+        model = ir.load(tmp_path / "model.onnx")
+        assert [value.name for value in model.graph.inputs] == list(
+            WeatherNextForecastTask.input_names
+        )
+        assert [value.name for value in model.graph.outputs] == list(
+            WeatherNextForecastTask.output_names
+        )
+        assert model.graph.outputs[0].shape == ir.Shape(["batch", 3, 4, 2])
+        assert model.graph.name == "weathernext_one_step_forecast"
+
+
+def test_weathernext_model_requires_weights_for_standard_build(tmp_path):
+    config = _config()
+    package = build_from_module(WeatherNextModel(config), config, task="weathernext-forecast")
+    package.apply_weights(create_demo_state_dict(config))
+    package.save(tmp_path, check_weights=True, progress_bar=False)
+
+
+def test_npz_inputs_infer_config(tmp_path):
+    input_path = tmp_path / "sample.npz"
+    np.savez(
+        input_path,
+        input_state=np.zeros((2, 5, 6, 3), dtype=np.float32),
+        forcings=np.zeros((2, 5, 6, 2), dtype=np.float32),
+        sample_noise=np.zeros((2, 5, 6, 1), dtype=np.float32),
+    )
+
+    feeds = load_npz_forecast_inputs(input_path)
+    config = infer_config_from_feeds(feeds, mesh_nodes=7, hidden_size=8)
+
+    assert config.lat == 5
+    assert config.lon == 6
+    assert config.mesh_nodes == 7
+    assert config.input_variables == 3
+    assert config.forcing_variables == 2
+    assert config.noise_channels == 1
+    assert config.output_variables == 3
diff --git a/src/mobius/tasks/__init__.py b/src/mobius/tasks/__init__.py
index d7f947ffc..2a780eb89 100644
--- a/src/mobius/tasks/__init__.py
+++ b/src/mobius/tasks/__init__.py
@@ -67,6 +67,7 @@
     "VAETask",
     "VideoDenoisingTask",
     "VisionLanguageTask",
+    "WeatherNextForecastTask",
     "WorldModelTask",
     "build_decoder_from_embeds",
     "build_embedding_from_features",
@@ -129,6 +130,7 @@
     QwenVLTask,
     VisionLanguageTask,
 )
+from mobius.tasks._weathernext import WeatherNextForecastTask
 from mobius.tasks._world_model import WorldModelTask
 
 # ---------------------------------------------------------------------------
@@ -181,6 +183,7 @@
     "ssm2-text-generation": SSM2CausalLMTask,
     "tts": TTSTask,
     "video-denoising": VideoDenoisingTask,
+    "weathernext-forecast": WeatherNextForecastTask,
     "world-model": WorldModelTask,
 }
 
diff --git a/src/mobius/tasks/_weathernext.py b/src/mobius/tasks/_weathernext.py
new file mode 100644
index 000000000..ab241e066
--- /dev/null
+++ b/src/mobius/tasks/_weathernext.py
@@ -0,0 +1,61 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""WeatherNext forecast task wiring."""
+
+from __future__ import annotations
+
+from typing import ClassVar
+
+import onnx_ir as ir
+from onnxscript import nn
+
+from mobius._configs import WeatherNextConfig
+from mobius._model_package import ModelPackage
+from mobius.tasks._base import ModelTask, _make_graph, _make_model
+
+
+class WeatherNextForecastTask(ModelTask):
+    """Build a one-step WeatherNext forecast graph.
+
+    Inputs:
+        - input_state: ``[batch, lat, lon, input_variables]``
+        - forcings: ``[batch, lat, lon, forcing_variables]``
+        - sample_noise: ``[batch, lat, lon, noise_channels]``
+
+    Outputs:
+        - next_state: ``[batch, lat, lon, output_variables]``
+    """
+
+    input_names: ClassVar[tuple[str, ...]] = ("input_state", "forcings", "sample_noise")
+    output_names: ClassVar[tuple[str, ...]] = ("next_state",)
+    model_roles: ClassVar[dict[str, str]] = {"model": "encoder"}
+
+    def build(
+        self,
+        module: nn.Module,
+        config: WeatherNextConfig,
+    ) -> ModelPackage:
+        config.validate()
+        batch = ir.SymbolicDim("batch")
+        graph, builder = _make_graph(name="weathernext_one_step_forecast")
+
+        input_state = builder.input(
+            self.input_names[0],
+            dtype=config.dtype,
+            shape=[batch, config.lat, config.lon, config.input_variables],
+        )
+        forcings = builder.input(
+            self.input_names[1],
+            dtype=config.dtype,
+            shape=[batch, config.lat, config.lon, config.forcing_variables],
+        )
+        sample_noise = builder.input(
+            self.input_names[2],
+            dtype=config.dtype,
+            shape=[batch, config.lat, config.lon, config.noise_channels],
+        )
+
+        next_state = module(builder.op, input_state, forcings, sample_noise)
+        builder.add_output(next_state, self.output_names[0])
+        return ModelPackage({"model": _make_model(graph)}, config=config)
diff --git a/tests/weathernext_example_test.py b/tests/weathernext_example_test.py
index 140f8d257..51543d773 100644
--- a/tests/weathernext_example_test.py
+++ b/tests/weathernext_example_test.py
@@ -3,11 +3,12 @@
 
 from __future__ import annotations
 
+import argparse
 import importlib.util
 import sys
 from pathlib import Path
 
-import onnx_ir as ir
+import numpy as np
 
 
 def _load_weathernext_example():
@@ -21,27 +22,35 @@ def _load_weathernext_example():
     return module
 
 
-def test_weathernext_demo_exports_one_step_forecast_model(tmp_path):
+def test_weathernext_example_infers_config_from_real_npz_inputs(tmp_path):
     example = _load_weathernext_example()
-    shape = example.WeatherNextDemoShape(
-        lat=3,
-        lon=4,
-        mesh_nodes=5,
-        input_variables=2,
-        forcing_variables=1,
+    input_path = tmp_path / "sample.npz"
+    np.savez(
+        input_path,
+        input_state=np.zeros((1, 3, 4, 2), dtype=np.float32),
+        forcings=np.zeros((1, 3, 4, 1), dtype=np.float32),
+        sample_noise=np.zeros((1, 3, 4, 1), dtype=np.float32),
+    )
+
+    args = argparse.Namespace(
+        input_data=str(input_path),
+        input_variable_names=None,
+        forcing_variable_names=None,
         noise_channels=1,
-        output_variables=2,
+        batch_index=0,
+        sample_noise_seed=0,
+        mesh_nodes=5,
         hidden_size=8,
+        intermediate_size=None,
+        num_hidden_layers=1,
+        dtype="f32",
     )
-
-    pkg = example.build_weathernext_demo_package(shape, dtype=ir.DataType.FLOAT)
-    pkg.save(tmp_path, check_weights=True, progress_bar=False)
-
-    model = ir.load(tmp_path / "model.onnx")
-    assert [value.name for value in model.graph.inputs] == [
-        "input_state",
-        "forcings",
-        "sample_noise",
-    ]
-    assert [value.name for value in model.graph.outputs] == ["next_state"]
-    assert model.graph.outputs[0].shape == ir.Shape(["batch", 3, 4, 2])
+    feeds = example._load_real_data(args)
+    config = example._config_from_args(args, feeds)
+
+    assert config.lat == 3
+    assert config.lon == 4
+    assert config.input_variables == 2
+    assert config.forcing_variables == 1
+    assert config.noise_channels == 1
+    assert config.output_variables == 2

From d7626adee9bde5ac26dc95a486782156f42743da Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Mon, 10 Aug 2026 13:58:42 +0000
Subject: [PATCH 10/18] Address WeatherNext review cleanup

Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
---
 src/mobius/integrations/weathernext.py | 5 ++---
 src/mobius/models/__init__.py          | 2 +-
 src/mobius/models/weathernext.py       | 6 +++++-
 3 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/src/mobius/integrations/weathernext.py b/src/mobius/integrations/weathernext.py
index 78fa1c096..26c2c8ea8 100644
--- a/src/mobius/integrations/weathernext.py
+++ b/src/mobius/integrations/weathernext.py
@@ -17,6 +17,8 @@
 from mobius._model_package import ModelPackage
 from mobius.models import WeatherNextModel
 
+_INPUT_NAMES = ("input_state", "forcings", "sample_noise")
+
 
 def load_npz_weights(path: str | Path) -> dict[str, torch.Tensor]:
     """Load a Mobius-aligned WeatherNext state dict from an ``.npz`` file."""
@@ -170,9 +172,6 @@ def infer_config_from_feeds(
     )
 
 
-_INPUT_NAMES = ("input_state", "forcings", "sample_noise")
-
-
 def _stack_xarray_variables(
     dataset: Any,
     names: list[str],
diff --git a/src/mobius/models/__init__.py b/src/mobius/models/__init__.py
index a1f6bb934..dad4c3f28 100644
--- a/src/mobius/models/__init__.py
+++ b/src/mobius/models/__init__.py
@@ -148,8 +148,8 @@
     "VideoAutoencoderModel",
     "Wav2Vec2ForCTCModel",
     "Wav2Vec2Model",
-    "WhisperForConditionalGeneration",
     "WeatherNextModel",
+    "WhisperForConditionalGeneration",
     "MLPWorldModel",
     "XLMCausalLMModel",
     "Zamba2CausalLMModel",
diff --git a/src/mobius/models/weathernext.py b/src/mobius/models/weathernext.py
index bf945b7de..84d1538fa 100644
--- a/src/mobius/models/weathernext.py
+++ b/src/mobius/models/weathernext.py
@@ -10,14 +10,18 @@
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 import numpy as np
 import onnx_ir as ir
-import torch
 from onnxscript import OpBuilder, nn
 
 from mobius._configs import WeatherNextConfig
 from mobius.components import Linear, get_activation
 
+if TYPE_CHECKING:
+    import torch
+
 
 def projection_matrix(rows: int, cols: int) -> np.ndarray:
     """Create a normalized deterministic projection between grid and mesh points."""

From 85c157c5d026d474f6c78c69e4c82cfd449a853e Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Mon, 10 Aug 2026 14:00:28 +0000
Subject: [PATCH 11/18] Harden WeatherNext data helpers

Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
---
 examples/weathernext.py                |  4 +-
 src/mobius/integrations/weathernext.py | 15 +++++-
 src/mobius/models/weathernext_test.py  | 19 +++++++-
 tests/weathernext_example_test.py      | 66 ++++++++++++--------------
 4 files changed, 66 insertions(+), 38 deletions(-)

diff --git a/examples/weathernext.py b/examples/weathernext.py
index e6af71f43..9c4ceb2e0 100644
--- a/examples/weathernext.py
+++ b/examples/weathernext.py
@@ -88,7 +88,9 @@ def _synthetic_feeds(config: WeatherNextConfig) -> dict[str, np.ndarray]:
     }
 
 
-def _config_from_args(args: argparse.Namespace, feeds: dict[str, np.ndarray] | None):
+def _config_from_args(
+    args: argparse.Namespace, feeds: dict[str, np.ndarray] | None
+) -> WeatherNextConfig:
     dtype = _resolve_dtype(args.dtype)
     if feeds is not None:
         return infer_config_from_feeds(
diff --git a/src/mobius/integrations/weathernext.py b/src/mobius/integrations/weathernext.py
index 26c2c8ea8..9d4bd22f9 100644
--- a/src/mobius/integrations/weathernext.py
+++ b/src/mobius/integrations/weathernext.py
@@ -80,6 +80,13 @@ def build_weathernext_package(
 def load_npz_forecast_inputs(path: str | Path) -> dict[str, np.ndarray]:
     """Load ``input_state``, ``forcings``, and ``sample_noise`` arrays from ``.npz``."""
     with np.load(path) as data:
+        missing = [name for name in _INPUT_NAMES if name not in data]
+        if missing:
+            expected = ", ".join(_INPUT_NAMES)
+            missing_names = ", ".join(missing)
+            raise ValueError(
+                f"WeatherNext input file is missing {missing_names}; expected keys: {expected}"
+            )
         feeds = {name: np.asarray(data[name], dtype=np.float32) for name in _INPUT_NAMES}
     return feeds
 
@@ -186,6 +193,12 @@ def _stack_xarray_variables(
         for dim in value.dims:
             if dim.lower() in {"time", "batch"}:
                 value = value.isel({dim: batch_index})
-        arrays.append(np.asarray(value, dtype=np.float32))
+        array = np.asarray(value, dtype=np.float32)
+        if array.ndim != 2:
+            raise ValueError(
+                f"Variable {name!r} must resolve to a 2-D lat/lon array after "
+                f"time/batch selection, got shape {array.shape}"
+            )
+        arrays.append(array)
     stacked = np.stack(arrays, axis=-1)
     return stacked[None, ...] if stacked.ndim == 3 else stacked
diff --git a/src/mobius/models/weathernext_test.py b/src/mobius/models/weathernext_test.py
index 908e93cff..63d51a9f4 100644
--- a/src/mobius/models/weathernext_test.py
+++ b/src/mobius/models/weathernext_test.py
@@ -7,7 +7,12 @@
 import onnx_ir as ir
 import pytest
 
-from mobius import WeatherNextConfig, WeatherNextForecastTask, WeatherNextModel, build_from_module
+from mobius import (
+    WeatherNextConfig,
+    WeatherNextForecastTask,
+    WeatherNextModel,
+    build_from_module,
+)
 from mobius.integrations.weathernext import (
     build_weathernext_package,
     create_demo_state_dict,
@@ -108,3 +113,15 @@ def test_npz_inputs_infer_config(tmp_path):
     assert config.forcing_variables == 2
     assert config.noise_channels == 1
     assert config.output_variables == 3
+
+
+def test_npz_inputs_report_missing_keys(tmp_path):
+    input_path = tmp_path / "missing.npz"
+    np.savez(
+        input_path,
+        input_state=np.zeros((1, 2, 3, 1), dtype=np.float32),
+        forcings=np.zeros((1, 2, 3, 1), dtype=np.float32),
+    )
+
+    with pytest.raises(ValueError, match="sample_noise"):
+        load_npz_forecast_inputs(input_path)
diff --git a/tests/weathernext_example_test.py b/tests/weathernext_example_test.py
index 51543d773..7a4809c21 100644
--- a/tests/weathernext_example_test.py
+++ b/tests/weathernext_example_test.py
@@ -3,28 +3,19 @@
 
 from __future__ import annotations
 
-import argparse
-import importlib.util
+import os
+import subprocess
 import sys
 from pathlib import Path
 
 import numpy as np
+import onnx_ir as ir
 
 
-def _load_weathernext_example():
-    path = Path(__file__).parents[1] / "examples" / "weathernext.py"
-    spec = importlib.util.spec_from_file_location("weathernext_example", path)
-    assert spec is not None
-    assert spec.loader is not None
-    module = importlib.util.module_from_spec(spec)
-    sys.modules[spec.name] = module
-    spec.loader.exec_module(module)
-    return module
-
-
-def test_weathernext_example_infers_config_from_real_npz_inputs(tmp_path):
-    example = _load_weathernext_example()
+def test_weathernext_example_builds_from_real_npz_inputs(tmp_path):
+    repo_root = Path(__file__).parents[1]
     input_path = tmp_path / "sample.npz"
+    output_dir = tmp_path / "weathernext"
     np.savez(
         input_path,
         input_state=np.zeros((1, 3, 4, 2), dtype=np.float32),
@@ -32,25 +23,30 @@ def test_weathernext_example_infers_config_from_real_npz_inputs(tmp_path):
         sample_noise=np.zeros((1, 3, 4, 1), dtype=np.float32),
     )
 
-    args = argparse.Namespace(
-        input_data=str(input_path),
-        input_variable_names=None,
-        forcing_variable_names=None,
-        noise_channels=1,
-        batch_index=0,
-        sample_noise_seed=0,
-        mesh_nodes=5,
-        hidden_size=8,
-        intermediate_size=None,
-        num_hidden_layers=1,
-        dtype="f32",
+    env = os.environ.copy()
+    env["PYTHONPATH"] = str(repo_root / "src")
+    subprocess.run(
+        [
+            sys.executable,
+            str(repo_root / "examples" / "weathernext.py"),
+            str(output_dir),
+            "--input-data",
+            str(input_path),
+            "--mesh-nodes",
+            "5",
+            "--hidden-size",
+            "8",
+        ],
+        check=True,
+        cwd=repo_root,
+        env=env,
     )
-    feeds = example._load_real_data(args)
-    config = example._config_from_args(args, feeds)
 
-    assert config.lat == 3
-    assert config.lon == 4
-    assert config.input_variables == 2
-    assert config.forcing_variables == 1
-    assert config.noise_channels == 1
-    assert config.output_variables == 2
+    model = ir.load(output_dir / "model.onnx")
+    assert [value.name for value in model.graph.inputs] == [
+        "input_state",
+        "forcings",
+        "sample_noise",
+    ]
+    assert [value.name for value in model.graph.outputs] == ["next_state"]
+    assert model.graph.outputs[0].shape == ir.Shape(["batch", 3, 4, 2])

From d9835c9a1fc457166970d4c033d7ff482ee5e699 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Mon, 10 Aug 2026 14:01:52 +0000
Subject: [PATCH 12/18] Simplify WeatherNext xarray stacking

Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
---
 src/mobius/integrations/weathernext.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/mobius/integrations/weathernext.py b/src/mobius/integrations/weathernext.py
index 9d4bd22f9..e9667ca1c 100644
--- a/src/mobius/integrations/weathernext.py
+++ b/src/mobius/integrations/weathernext.py
@@ -185,6 +185,8 @@ def _stack_xarray_variables(
     *,
     batch_index: int,
 ) -> np.ndarray:
+    if not names:
+        raise ValueError("At least one xarray variable name is required")
     arrays = []
     for name in names:
         if name not in dataset:
@@ -201,4 +203,4 @@ def _stack_xarray_variables(
             )
         arrays.append(array)
     stacked = np.stack(arrays, axis=-1)
-    return stacked[None, ...] if stacked.ndim == 3 else stacked
+    return stacked[None, ...]

From 460f50b94f3370f519826cecd1e140acc90e1ddb Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Mon, 10 Aug 2026 14:03:12 +0000
Subject: [PATCH 13/18] Finalize WeatherNext review fixes

Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
---
 src/mobius/integrations/weathernext.py | 19 +++++++++----------
 src/mobius/tasks/_weathernext.py       |  2 +-
 2 files changed, 10 insertions(+), 11 deletions(-)

diff --git a/src/mobius/integrations/weathernext.py b/src/mobius/integrations/weathernext.py
index e9667ca1c..8ed2a11dc 100644
--- a/src/mobius/integrations/weathernext.py
+++ b/src/mobius/integrations/weathernext.py
@@ -128,19 +128,18 @@ def load_xarray_forecast_inputs(
             forcing_variables,
             batch_index=batch_index,
         )
+        rng = np.random.default_rng(sample_noise_seed)
+        sample_noise = rng.standard_normal(
+            (input_state.shape[0], input_state.shape[1], input_state.shape[2], noise_channels)
+        ).astype(np.float32)
+        return {
+            "input_state": input_state,
+            "forcings": forcings,
+            "sample_noise": sample_noise,
+        }
     finally:
         dataset.close()
 
-    rng = np.random.default_rng(sample_noise_seed)
-    sample_noise = rng.standard_normal(
-        (input_state.shape[0], input_state.shape[1], input_state.shape[2], noise_channels)
-    ).astype(np.float32)
-    return {
-        "input_state": input_state,
-        "forcings": forcings,
-        "sample_noise": sample_noise,
-    }
-
 
 def infer_config_from_feeds(
     feeds: dict[str, np.ndarray],
diff --git a/src/mobius/tasks/_weathernext.py b/src/mobius/tasks/_weathernext.py
index ab241e066..719bfe550 100644
--- a/src/mobius/tasks/_weathernext.py
+++ b/src/mobius/tasks/_weathernext.py
@@ -29,7 +29,7 @@ class WeatherNextForecastTask(ModelTask):
 
     input_names: ClassVar[tuple[str, ...]] = ("input_state", "forcings", "sample_noise")
     output_names: ClassVar[tuple[str, ...]] = ("next_state",)
-    model_roles: ClassVar[dict[str, str]] = {"model": "encoder"}
+    model_roles: ClassVar[dict[str, str]] = {"model": "forecast"}
 
     def build(
         self,

From f99b0dc60a228fc03b85a22168254deaec1eea84 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Mon, 10 Aug 2026 14:04:31 +0000
Subject: [PATCH 14/18] Expose WeatherNext output variable count

Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
---
 examples/weathernext.py                | 9 +++++++--
 src/mobius/integrations/weathernext.py | 5 ++++-
 src/mobius/models/weathernext_test.py  | 9 +++++++--
 3 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/examples/weathernext.py b/examples/weathernext.py
index 9c4ceb2e0..3018a466d 100644
--- a/examples/weathernext.py
+++ b/examples/weathernext.py
@@ -97,6 +97,7 @@ def _config_from_args(
             feeds,
             mesh_nodes=args.mesh_nodes,
             hidden_size=args.hidden_size,
+            output_variables=args.output_variables,
             intermediate_size=args.intermediate_size,
             num_hidden_layers=args.num_hidden_layers,
             dtype=dtype,
@@ -108,7 +109,9 @@ def _config_from_args(
         input_variables=args.input_variables,
         forcing_variables=args.forcing_variables,
         noise_channels=args.noise_channels,
-        output_variables=args.output_variables,
+        output_variables=args.input_variables
+        if args.output_variables is None
+        else args.output_variables,
         hidden_size=args.hidden_size,
         intermediate_size=args.intermediate_size or 4 * args.hidden_size,
         num_hidden_layers=args.num_hidden_layers,
@@ -148,7 +151,9 @@ def _parse_args() -> argparse.Namespace:
         "--noise-channels", type=int, default=2, help="Stochastic noise channels."
     )
     parser.add_argument(
-        "--output-variables", type=int, default=5, help="Synthetic output channels."
+        "--output-variables",
+        type=int,
+        help="Output channels. Defaults to input-variable count for synthetic and real data.",
     )
     parser.add_argument("--hidden-size", type=int, default=16, help="Latent feature size.")
     parser.add_argument("--intermediate-size", type=int, help="Mesh MLP intermediate size.")
diff --git a/src/mobius/integrations/weathernext.py b/src/mobius/integrations/weathernext.py
index 8ed2a11dc..ed94d29fa 100644
--- a/src/mobius/integrations/weathernext.py
+++ b/src/mobius/integrations/weathernext.py
@@ -146,6 +146,7 @@ def infer_config_from_feeds(
     *,
     mesh_nodes: int,
     hidden_size: int,
+    output_variables: int | None = None,
     intermediate_size: int | None = None,
     num_hidden_layers: int = 1,
     dtype: ir.DataType = ir.DataType.FLOAT,
@@ -170,7 +171,9 @@ def infer_config_from_feeds(
         input_variables=int(input_state.shape[3]),
         forcing_variables=int(forcings.shape[3]),
         noise_channels=int(sample_noise.shape[3]),
-        output_variables=int(input_state.shape[3]),
+        output_variables=int(
+            input_state.shape[3] if output_variables is None else output_variables
+        ),
         hidden_size=hidden_size,
         intermediate_size=intermediate_size or 4 * hidden_size,
         num_hidden_layers=num_hidden_layers,
diff --git a/src/mobius/models/weathernext_test.py b/src/mobius/models/weathernext_test.py
index 63d51a9f4..982022eec 100644
--- a/src/mobius/models/weathernext_test.py
+++ b/src/mobius/models/weathernext_test.py
@@ -104,7 +104,12 @@ def test_npz_inputs_infer_config(tmp_path):
     )
 
     feeds = load_npz_forecast_inputs(input_path)
-    config = infer_config_from_feeds(feeds, mesh_nodes=7, hidden_size=8)
+    config = infer_config_from_feeds(
+        feeds,
+        mesh_nodes=7,
+        hidden_size=8,
+        output_variables=4,
+    )
 
     assert config.lat == 5
     assert config.lon == 6
@@ -112,7 +117,7 @@ def test_npz_inputs_infer_config(tmp_path):
     assert config.input_variables == 3
     assert config.forcing_variables == 2
     assert config.noise_channels == 1
-    assert config.output_variables == 3
+    assert config.output_variables == 4
 
 
 def test_npz_inputs_report_missing_keys(tmp_path):

From 19c6ae22633c06837a7c7f1c37b64a16d9d4f9e4 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Mon, 10 Aug 2026 14:05:42 +0000
Subject: [PATCH 15/18] Clarify WeatherNext data docs

Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
---
 src/mobius/integrations/weathernext.py | 4 +++-
 src/mobius/models/weathernext.py       | 4 ++--
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/src/mobius/integrations/weathernext.py b/src/mobius/integrations/weathernext.py
index ed94d29fa..cb9210a78 100644
--- a/src/mobius/integrations/weathernext.py
+++ b/src/mobius/integrations/weathernext.py
@@ -104,7 +104,9 @@ def load_xarray_forecast_inputs(
 
     The selected variables must share latitude/longitude dimensions. Optional
     time dimensions are indexed by ``batch_index`` and each variable is stacked
-    into the trailing channel dimension expected by the ONNX graph.
+    into the trailing channel dimension expected by the ONNX graph. The default
+    noise seed is deterministic so examples are reproducible; pass a different
+    seed when sampling stochastic forecast noise for real workflows.
     """
     try:
         import xarray as xr
diff --git a/src/mobius/models/weathernext.py b/src/mobius/models/weathernext.py
index 84d1538fa..ea99d36fa 100644
--- a/src/mobius/models/weathernext.py
+++ b/src/mobius/models/weathernext.py
@@ -94,8 +94,7 @@ def forward(
         # noise: [B, lat, lon, input+forcing+noise].
         grid_features = op.Concat(input_state, forcings, sample_noise, axis=-1)
 
-        # Encode each lat/lon cell independently, then flatten the spatial grid:
-        # [B, lat, lon, hidden] -> [B, grid_points, hidden].
+        # Encode each lat/lon cell independently: [B, lat, lon, hidden].
         grid_latent = self._activation(op, self.grid_encoder(op, grid_features))
         batch_dim = op.Shape(grid_latent, start=0, end=1)
         flat_grid_shape = op.Concat(
@@ -103,6 +102,7 @@ def forward(
             op.Constant(value_ints=[config.grid_points, config.hidden_size]),
             axis=0,
         )
+        # Flatten the spatial grid: [B, lat, lon, hidden] -> [B, grid_points, hidden].
         grid_points = op.Reshape(grid_latent, flat_grid_shape)
 
         # Aggregate encoded grid cells onto mesh nodes:

From f588c1ae11472c843ade9aa1941bc9e2aeba590e Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Mon, 10 Aug 2026 14:07:15 +0000
Subject: [PATCH 16/18] Make WeatherNext projections explicit

Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
---
 src/mobius/integrations/weathernext.py | 1 +
 src/mobius/models/weathernext.py       | 8 ++++----
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/src/mobius/integrations/weathernext.py b/src/mobius/integrations/weathernext.py
index cb9210a78..52a440a41 100644
--- a/src/mobius/integrations/weathernext.py
+++ b/src/mobius/integrations/weathernext.py
@@ -130,6 +130,7 @@ def load_xarray_forecast_inputs(
             forcing_variables,
             batch_index=batch_index,
         )
+        # The seed controls only the generated stochastic noise input.
         rng = np.random.default_rng(sample_noise_seed)
         sample_noise = rng.standard_normal(
             (input_state.shape[0], input_state.shape[1], input_state.shape[2], noise_channels)
diff --git a/src/mobius/models/weathernext.py b/src/mobius/models/weathernext.py
index ea99d36fa..78a6e971e 100644
--- a/src/mobius/models/weathernext.py
+++ b/src/mobius/models/weathernext.py
@@ -106,9 +106,9 @@ def forward(
         grid_points = op.Reshape(grid_latent, flat_grid_shape)
 
         # Aggregate encoded grid cells onto mesh nodes:
-        # [mesh_nodes, grid_points] @ [B, grid_points, hidden]
+        # [1, mesh_nodes, grid_points] @ [B, grid_points, hidden]
         # -> [B, mesh_nodes, hidden].
-        mesh_latent = op.MatMul(self.grid_to_mesh, grid_points)
+        mesh_latent = op.MatMul(op.Unsqueeze(self.grid_to_mesh, [0]), grid_points)
 
         # Apply one or more residual mesh-update MLP blocks.
         for update_in, update_out in zip(
@@ -118,9 +118,9 @@ def forward(
             mesh_latent = op.Add(mesh_latent, update_out(op, mesh_delta))
 
         # Decode mesh latents back to grid cells and retain the encoded-grid
-        # residual: [grid_points, mesh_nodes] @ [B, mesh_nodes, hidden]
+        # residual: [1, grid_points, mesh_nodes] @ [B, mesh_nodes, hidden]
         # -> [B, grid_points, hidden].
-        grid_delta = op.MatMul(self.mesh_to_grid, mesh_latent)
+        grid_delta = op.MatMul(op.Unsqueeze(self.mesh_to_grid, [0]), mesh_latent)
         grid_points = op.Add(grid_points, grid_delta)
 
         # Return one forecast step on the original grid:

From 6156702127cc813d88e59490c01b2face5998992 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Tue, 11 Aug 2026 15:29:41 +0000
Subject: [PATCH 17/18] Cast WeatherNext feeds for f16 inference

Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
---
 examples/weathernext.py           | 25 +++++++++++++++---
 tests/weathernext_example_test.py | 42 +++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+), 3 deletions(-)

diff --git a/examples/weathernext.py b/examples/weathernext.py
index 3018a466d..ec202304f 100644
--- a/examples/weathernext.py
+++ b/examples/weathernext.py
@@ -88,6 +88,21 @@ def _synthetic_feeds(config: WeatherNextConfig) -> dict[str, np.ndarray]:
     }
 
 
+def _numpy_dtype(dtype: ir.DataType) -> type[np.float32] | type[np.float16]:
+    if dtype == ir.DataType.FLOAT:
+        return np.float32
+    if dtype == ir.DataType.FLOAT16:
+        return np.float16
+    raise ValueError(f"Unsupported WeatherNext feed dtype: {dtype}")
+
+
+def _cast_feeds_to_dtype(
+    feeds: dict[str, np.ndarray], dtype: ir.DataType
+) -> dict[str, np.ndarray]:
+    feed_dtype = _numpy_dtype(dtype)
+    return {name: np.asarray(value, dtype=feed_dtype) for name, value in feeds.items()}
+
+
 def _config_from_args(
     args: argparse.Namespace, feeds: dict[str, np.ndarray] | None
 ) -> WeatherNextConfig:
@@ -119,7 +134,9 @@ def _config_from_args(
     )
 
 
-def _run_with_ort(output_dir: str, feeds: dict[str, np.ndarray]) -> None:
+def _run_with_ort(
+    output_dir: str, feeds: dict[str, np.ndarray], dtype: ir.DataType
+) -> None:
     try:
         import onnxruntime as ort
     except ImportError:
@@ -128,7 +145,7 @@ def _run_with_ort(output_dir: str, feeds: dict[str, np.ndarray]) -> None:
 
     model_path = os.path.join(output_dir, "model.onnx")
     sess = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
-    (next_state,) = sess.run(["next_state"], feeds)
+    (next_state,) = sess.run(["next_state"], _cast_feeds_to_dtype(feeds, dtype))
     print(f"Inference output next_state shape: {next_state.shape}")
     print(f"Inference output range: [{next_state.min():.6f}, {next_state.max():.6f}]")
 
@@ -216,7 +233,9 @@ def main() -> None:
 
     if args.run or args.validate:
         _run_with_ort(
-            args.output_dir, feeds if feeds is not None else _synthetic_feeds(config)
+            args.output_dir,
+            feeds if feeds is not None else _synthetic_feeds(config),
+            config.dtype,
         )
 
 
diff --git a/tests/weathernext_example_test.py b/tests/weathernext_example_test.py
index 7a4809c21..c141327fa 100644
--- a/tests/weathernext_example_test.py
+++ b/tests/weathernext_example_test.py
@@ -10,6 +10,7 @@
 
 import numpy as np
 import onnx_ir as ir
+import pytest
 
 
 def test_weathernext_example_builds_from_real_npz_inputs(tmp_path):
@@ -50,3 +51,44 @@ def test_weathernext_example_builds_from_real_npz_inputs(tmp_path):
     ]
     assert [value.name for value in model.graph.outputs] == ["next_state"]
     assert model.graph.outputs[0].shape == ir.Shape(["batch", 3, 4, 2])
+
+
+def test_weathernext_example_runs_f16_real_npz_inputs(tmp_path):
+    pytest.importorskip("onnxruntime")
+
+    repo_root = Path(__file__).parents[1]
+    input_path = tmp_path / "sample.npz"
+    output_dir = tmp_path / "weathernext-f16"
+    np.savez(
+        input_path,
+        input_state=np.zeros((1, 3, 4, 2), dtype=np.float32),
+        forcings=np.zeros((1, 3, 4, 1), dtype=np.float32),
+        sample_noise=np.zeros((1, 3, 4, 1), dtype=np.float32),
+    )
+
+    env = os.environ.copy()
+    env["PYTHONPATH"] = str(repo_root / "src")
+    result = subprocess.run(
+        [
+            sys.executable,
+            str(repo_root / "examples" / "weathernext.py"),
+            str(output_dir),
+            "--input-data",
+            str(input_path),
+            "--mesh-nodes",
+            "5",
+            "--hidden-size",
+            "8",
+            "--dtype",
+            "f16",
+            "--run",
+            "--validate",
+        ],
+        check=True,
+        cwd=repo_root,
+        env=env,
+        text=True,
+        capture_output=True,
+    )
+
+    assert "Inference output next_state shape: (1, 3, 4, 2)" in result.stdout

From f9191fdfe73bda22850217b77d8967378da0789b Mon Sep 17 00:00:00 2001
From: justinchuby <justinchuby@users.noreply.github.com>
Date: Tue, 11 Aug 2026 15:46:46 +0000
Subject: [PATCH 18/18] Fix WeatherNext inference lint

Apply current formatting and typing rules to the dtype-aware ONNX Runtime feed helper.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: justinchuby <justinchuby@users.noreply.github.com>
---
 examples/weathernext.py | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/examples/weathernext.py b/examples/weathernext.py
index ec202304f..7948dae56 100644
--- a/examples/weathernext.py
+++ b/examples/weathernext.py
@@ -88,7 +88,7 @@ def _synthetic_feeds(config: WeatherNextConfig) -> dict[str, np.ndarray]:
     }
 
 
-def _numpy_dtype(dtype: ir.DataType) -> type[np.float32] | type[np.float16]:
+def _numpy_dtype(dtype: ir.DataType) -> type[np.float32 | np.float16]:
     if dtype == ir.DataType.FLOAT:
         return np.float32
     if dtype == ir.DataType.FLOAT16:
@@ -134,9 +134,7 @@ def _config_from_args(
     )
 
 
-def _run_with_ort(
-    output_dir: str, feeds: dict[str, np.ndarray], dtype: ir.DataType
-) -> None:
+def _run_with_ort(output_dir: str, feeds: dict[str, np.ndarray], dtype: ir.DataType) -> None:
     try:
         import onnxruntime as ort
     except ImportError: