[alxmrs]

A draft, opened early so the work is visible. This is demo 4 in benchmarks/ml: Samudra 2's network written as SQL views, checked against Samudra's own PyTorch modules. Samudra is an emulator of the global ocean. Its network is a ConvNeXt U-Net: convolutions, batch norm, a capped GELU, average pooling, and zonally periodic bilinear upsampling. Toward #291 (a U-Net demo for einfold to optimize convolutions against).

What's here

benchmarks/ml/04_samudra.py: each of Samudra's modules is a SQL function named after it, and each ConvNeXt block is a view over a state table (c, lat, lon, v).

• conv: each input cell, through each tap of the kernel, reaches the output cell the tap points to. Longitude wraps around the globe (modulo the grid's width), and latitude stops at the poles, which is Samudra's zero padding.
• batch_norm: the module's running statistics.
• capped_gelu: GELU capped at 10, with erf computed by Abramowitz and Stegun's formula 7.1.26 (error below 1.5e-7), since neither DataFusion nor DuckDB has erf.
• avg_pool: the mean of each 2 × 2 block of cells.
• upsample: each fine cell reads four coarse cells with bilinear weights, then adds the skip connection.

The check builds Samudra's UNetBackbone, ConvNeXtBlock and Samudra modules, as Samudra 2's configuration does, from m2lines/Samudra@4f86be4. They get random weights and batch-norm statistics at reduced width (4 to 10 channels), on a small grid, and run in float32. --levels sets the depth, from 1 to 4.

First result, from the prototype at one level on a 10 × 18 grid: the forecast matches Samudra's PyTorch to 1.75e-7 relative on DataFusion, DuckDB and einfold. That's the error of the erf formula.

Found on the way

• DataFusion 55: eliminate_cross_join drops a join key that reads two tables of a cross join, and returns wrong rows — draft for upstream #299: a wrong answer in DataFusion 55. eliminate_cross_join drops a join key that reads two tables of a cross join. The first spelling of the upsampling hit it. The upsampling now reads its source cells from a CTE, which is also the clearer spelling.
• Portability:
  • DuckDB's / on integers returns a float, and DataFusion has no //, so pooling writes floor(lat / 2.0).
  • DataFusion has signum and DuckDB has sign, so GELU uses CASE.
  • A dotted table name (decoder.weight) is read as a schema and a table, so parameter tables use underscores.
• Samudra's model is hard to import on its own: samudra.models pulls in Microsoft's Aurora, timm, torchvision and the data pipeline. This draft installs Samudra with --no-deps plus those packages.

Not done yet

• Packaging. The script's samudra @ git+… dependency pulls in all of Samudra's dependencies (numpy<2, skypilot, wandb, and others). CI needs a CPU build of PyTorch. tests/test_ml_demos.py lists the demo with no run until then.
• Results at 2 to 4 levels, and plan sizes. Each block reads its input twice, and DataFusion plans a view again at every read (DataFusion plans a CTE once per reference, so ML queries grow exponentially with depth #267), so the plan should grow quickly with depth.
• The full run: Samudra 2's released 1° weights (M2LInES/Samudra2 on Hugging Face) on the public OM4 1° data (OSN), with Samudra's normalization and channel order.
• The README row, the demo's docs, and numbers on Tracking: ML in SQL as fast as JAX on CPU #285.
• A follow-up PR: ddx's gradients, with respect to the input and to every weight table.

🤖 Generated with Claude Code

https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

Share

• X
• Mastodon
• Reddit
• LinkedIn

[alxmrs]

🌀 Claude: The pushed script, run at 1 and 2 levels on a 10 × 20 grid. Measured on an i7-8700 with datafusion 55.0.0, DuckDB 1.5.6, einfold 0.2.0rc2 and torch 2.15.0 (CPU).

levels	host	time	error against Samudra's PyTorch	operators
1	DuckDB	0.59 s	2e-07	280
1	DataFusion	0.25 s	2e-07	599
1	einfold	0.33 s	2e-07	1
2	DuckDB	308 s	3e-07	693
2	DataFusion	1.65 s	3e-07	3,413
2	einfold	1.99 s	3e-07	1

• Every host matches Samudra's modules to the error of the erf formula. That covers the skip connection, the odd-sized padding after upsampling, and the final block's dilation.
• DataFusion's plan grows 5.7× with one more level (599 to 3,413 operators). Each ConvNeXt block reads its input twice, and DataFusion plans a view again at every read (DataFusion plans a CTE once per reference, so ML queries grow exponentially with depth #267).
• DuckDB's 308 s is almost certainly planning, not computing. The grid has 200 cells, and DuckDB's planning time grows with each view in the same way on GraphCast's views (DuckDB plans views that share reads in time exponential in their depth — draft for upstream #292). I haven't measured it separately yet.
• einfold shows one operator because its plan crosses the FFI as one node (Four gaps for a query planner in datafusion-python's process, through its FFI — draft for upstream einfold#58). It doesn't fuse convolutions yet, which is ML demo: a U-Net in SQL, for einfold to optimize convolutions against #291's purpose.