ilnmtlbnm/jacobi2000-model20

jacobi2000 model20

A small transformer that steps a 2D shallow-water simulation on a 64x64 grid, trained in PyTorch from a reference solver. Patch size 4, 4 layers, 4 heads, width 128, attention restricted to a 3x3 stencil of patches.

Code, training and the inference engine (Rust, CPU and WebGPU): github.com/ilnmtlbnm/jacobi2000

Run it

hf download ilnmtlbnm/jacobi2000-model20 --local-dir models/model20
cargo run -p engine-cli --release --bin rollout -- \
    --model models/model20 --steps 50 --backend cpu

config.json holds the architecture, the input normalisation and the physical constants of the training data (dx, dt, g, depth, nu).