Skip to content

Add WeatherNext-style ONNX conversion demo - #469

Draft
justinchuby with Copilot wants to merge 18 commits into
mainfrom
copilot/convert-weather-next
Draft

Add WeatherNext-style ONNX conversion demo#469
justinchuby with Copilot wants to merge 18 commits into
mainfrom
copilot/convert-weather-next

Conversation

Copilot AI commented Aug 10, 2026

Copy link
Copy Markdown
Contributor

WeatherNext uses a JAX/Haiku + xarray architecture that does not map to the existing HuggingFace model registry flow. This PR adds a runnable ONNX demo that sketches the WeatherNext one-step forecast contract and graph workflow under examples/.

  • Demo workflow

    • Adds examples/weathernext.py.
    • Builds a small deterministic grid→mesh→grid forecast graph.
    • Saves the result as a standard Mobius ModelPackage.
    • Supports optional ONNX Runtime validation with --validate.
  • WeatherNext-style graph contract

    • Inputs:
      • input_state: weather variables on a lat/lon grid
      • forcings: known future forcing variables
      • sample_noise: stochastic FGN-style noise channels
    • Output:
      • next_state: one-step forecast grid
  • Coverage

    • Adds tests/weathernext_example_test.py to verify the exported ONNX graph I/O contract.

Example:

PYTHONPATH=src python examples/weathernext.py output/weathernext-mini --validate

@CLAassistant

CLAassistant commented Aug 10, 2026

Copy link
Copy Markdown

CLA assistant check
Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you all sign our Contributor License Agreement before we can accept your contribution.
1 out of 2 committers have signed the CLA.

✅ justinchuby
❌ Copilot
You have signed the CLA already but the status is still pending? Let us recheck it.

Copilot AI linked an issue Aug 10, 2026 that may be closed by this pull request
Copilot AI and others added 7 commits August 10, 2026 06:22
Signed-off-by: GitHub <noreply@github.com>

Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Signed-off-by: GitHub <noreply@github.com>

Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Signed-off-by: GitHub <noreply@github.com>

Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Signed-off-by: GitHub <noreply@github.com>

Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Signed-off-by: GitHub <noreply@github.com>

Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Signed-off-by: GitHub <noreply@github.com>

Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Signed-off-by: GitHub <noreply@github.com>

Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Comment thread examples/weathernext.py Fixed
Copilot AI changed the title [WIP] Explore conversion of weather next to onnx Add WeatherNext-style ONNX conversion demo Aug 10, 2026
Copilot AI requested a review from justinchuby August 10, 2026 06:33
@justinchuby

Copy link
Copy Markdown
Member

@copilot what I wanted was a formal support of this model in the modeling/integation files, and a demo script that can run this model on real data.

Copilot AI and others added 8 commits August 10, 2026 13:57
Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>

Copilot AI commented Aug 10, 2026

Copy link
Copy Markdown
Contributor Author

@copilot what I wanted was a formal support of this model in the modeling/integation files, and a demo script that can run this model on real data.

Implemented in f588c1a (with the WeatherNext support commits immediately before it): the WeatherNext graph is now formalized under mobius config/model/task/integration modules, and examples/weathernext.py uses that path with NPZ/xarray real-data inputs plus optional converted weights.

@justinchuby justinchuby left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The advertised f16 validation path is currently not executable.

Comment thread examples/weathernext.py
Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
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>
@justinchuby
justinchuby dismissed their stale review August 11, 2026 15:46

The requested f16 inference fix and regression test landed in 6156702; follow-up lint cleanup landed in f9191fd.

@github-actions

Copy link
Copy Markdown

Performance Comparison

Comparing 78437cdf9191fd

Model Metric Baseline Current Delta
bert (feature-extraction) model_size_bytes 359 KB 359 KB +0.0%
bert (feature-extraction) num_nodes 60 60 +0.0%
falcon model_size_bytes 364 KB 364 KB +0.0%
falcon num_nodes 68 68 +0.0%
gemma2 model_size_bytes 428 KB 428 KB +0.0%
gemma2 num_nodes 107 107 +0.0%
gpt2 model_size_bytes 388 KB 388 KB +0.0%
gpt2 num_nodes 54 54 +0.0%
llama model_size_bytes 425 KB 425 KB +0.0%
llama num_nodes 62 62 +0.0%
llama (static-cache) model_size_bytes 425 KB 425 KB +0.0%
llama (static-cache) num_nodes 58 58 +0.0%
mamba (ssm-text-generation) model_size_bytes 296 KB 296 KB +0.0%
mamba (ssm-text-generation) num_nodes 98 98 +0.0%
phi3 model_size_bytes 421 KB 421 KB +0.0%
phi3 num_nodes 60 60 +0.0%
phi3 (static-cache) model_size_bytes 421 KB 421 KB +0.0%
phi3 (static-cache) num_nodes 56 56 +0.0%
qwen2 model_size_bytes 425 KB 425 KB +0.0%
qwen2 num_nodes 62 62 +0.0%
qwen2 (static-cache) model_size_bytes 425 KB 425 KB +0.0%
qwen2 (static-cache) num_nodes 58 58 +0.0%
qwen3_5_moe (hybrid-text-generation) model_size_bytes 506 KB 506 KB +0.0%
qwen3_5_moe (hybrid-text-generation) num_nodes 275 275 +0.0%
qwen3_5_text (hybrid-text-generation) model_size_bytes 458 KB 458 KB +0.0%
qwen3_5_text (hybrid-text-generation) num_nodes 129 129 +0.0%
qwen3_5_vl (hybrid-qwen-vl) model_size_bytes 977 KB 977 KB +0.0%
qwen3_5_vl (hybrid-qwen-vl) num_nodes 431 431 +0.0%
t5 (seq2seq) model_size_bytes 836 KB 836 KB +0.0%
t5 (seq2seq) num_nodes 166 166 +0.0%
whisper (speech-to-text) model_size_bytes 1008 KB 1008 KB +0.0%
whisper (speech-to-text) num_nodes 128 128 +0.0%

No performance regressions.

@github-actions

Copy link
Copy Markdown

🏗️ Architecture Diff

Comparing 78437cdf9191fd

Model Sub-model Changes Status
bert (feature-extraction) model 0
falcon model 0
gemma2 model 0
gemma4 (gemma4) decoder 0
gemma4 (gemma4) embedding 0
gemma4 (gemma4) vision_encoder 0
gemma4_text model 0
gpt2 model 0
llama model 0
llama (static-cache) model 0
mamba (ssm-text-generation) model 0
phi3 model 0
phi3 (static-cache) model 0
qwen model 0
qwen (static-cache) model 0
qwen2 model 0
qwen2 (static-cache) model 0
qwen2_moe model 0
qwen2_moe (static-cache) model 0
qwen3 model 0
qwen3 (static-cache) model 0
qwen3_5_moe (hybrid-text-generation) model 0
qwen3_5_text (hybrid-text-generation) model 0
qwen3_5_vl (hybrid-qwen-vl) decoder 0
qwen3_5_vl (hybrid-qwen-vl) embedding 0
qwen3_5_vl (hybrid-qwen-vl) vision_encoder 0
qwen3_moe model 0
qwen3_moe (static-cache) model 0
qwen3_next (hybrid-text-generation) model 0
t5 (seq2seq) decoder 0
t5 (seq2seq) encoder 0
whisper (speech-to-text) decoder 0
whisper (speech-to-text) encoder 0

No architecture changes detected.


Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Convert weather next

3 participants