[asjad2401]

Summary

Adds backend="wgpu" — the RFC's Stage 2 promise. All ops now run on Burn's burn-dispatch backend, which forwards each op to the backend selected by the tensor's device, so the interpreter stays backend-agnostic without making every op generic.

burn.available_backends()                           # ['flex', 'wgpu']
model = burn.load_onnx("model.onnx", backend="wgpu")

ResNet-18 (torchvision export, Apple M-series, median of 15)

backend	batch 1	batch 8
burn-python wgpu (Metal GPU)	6.1 ms	35.6 ms
burn-python flex (CPU)	22.4 ms	135.7 ms
ONNX Runtime (CPU)	14.9 ms	123.0 ms

All four ResNet-18 exports match ORT on wgpu too (max logit diff ~1e-5, same top-1).

Design

• type B = Dispatch; OnnxModel holds its DispatchDevice. Ops take the device from their inputs (B::float_device); only weight loading, inputs and inline constants need it passed in.
• Native shading language per target. burn-dispatch allows one wgpu flavor per build, so burn-dispatch is declared per target: metal on macOS, vulkan elsewhere. The portable webgpu (WGSL) path was measured too — ~4× slower (batch 8: 128 ms vs 32 ms) since WGSL can't use matrix hardware. wgpu loads the driver at runtime, so CPU-only machines still work with "flex".
• The unused burn umbrella crate is replaced by burn-dispatch.

Fixes found by running on a stricter backend

• Rank broadcasting. ONNX allows [N,C] + [C], Gemm with rank-1 C, [N,C,H,W] * [C,1,1]. Burn needs equal ranks; Flex tolerated the mismatch, wgpu doesn't. New ops/broadcast.rs aligns ranks numpy-style for elementwise ops and Linear/Gemm bias.
• Gemm alpha/beta and BN epsilon use float_*_scalar instead of 1-element tensors.
• Worker-thread failures. With no usable GPU, cubecl panics on its device thread and we only saw RecvError. New src/panic.rs records the first panic on any thread; model loading is now guarded too:
  RuntimeError: failed to load model on backend 'wgpu': No Discrete GPU device found, adapters [], ...

Testing

• BURN_BACKEND=wgpu|flex selects the backend for test_ops.py / test_resnet.py. Locally (Metal): 49/49 op checks and 16/16 ResNet checks pass on both backends; test_bridge.py, compare_ort.py pass.
• New op tests: cross-rank broadcasting, Gemm alpha/beta with rank-1 C, unknown backend → ValueError.
• New CI job wgpu tests (software Vulkan) runs the op tests through Mesa's lavapipe — checks the GPU code path on hosted runners (not its speed). Not in required checks.

🤖 Generated with Claude Code