Add an auto mode for inference dtype

[wrignj08]

inference_dtype defaults to torch.float32, and on hardware with a reduced-precision path that leaves throughput on the table — a Perth Sentinel-2 tile ran model inference in 27.5s at float32 against 21.1s at bfloat16, identical output to six decimal places. The examples pass inference_dtype="bf16" explicitly; anyone copying the README does not.

Pass inference_dtype="auto" and the device is measured rather than assumed.

The default is unchanged, so this changes nothing until it is asked for.

Why measuring, not a capability check

A device will accept bfloat16 tensors and run every kernel on them while emulating the arithmetic in software — far slower than the float32 it was meant to beat. On this machine's CPU an emulated convolution is ~100× slower. A capability flag is not evidence of a speedup.

What the measurement decides, and what it doesn't

Whether to drop precision at all is decided by measurement. That signal is large and consistent: ~100× on an emulating CPU, 1.5–5× on MPS.

Which dtype to drop to is decided by preference, not measurement. A single convolution turned out not to predict how a whole model ranks two dtypes sharing one hardware path — the probe put float16 ahead of bfloat16 here (0.6 ms against 1.0–2.0 ms) while the real model ranked them the other way (21.1s against 23.3s), and bfloat16's probe figure moved by 2× between runs. Letting it cast that vote returned float16, bfloat16 and float32 on alternate runs of the same machine. bfloat16 now wins on numerics — it keeps float32's exponent range where float16 overflows to inf — and float16 is used only where bfloat16 will not run.

Two details the timing has to get right

• Repeat until the window is long enough. Each measurement repeats the convolution until it has run 50 ms, then divides. A single convolution is a few milliseconds on a GPU, where scheduling jitter is a large fraction of the total. Repeating is much cheaper than enlarging the tensor, which grows memory quadratically.
• Screen before measuring. Candidates get one small call first, and anything already slower than float32 is dropped. Sizing the real probe from float32 and then running an emulated candidate at that size turned a sub-second selection into 212 s on CPU.

A candidate must beat float32 by 10% to be used at all; below that the difference is noise and float32 is what the models were trained in. The result is cached per device, so a batch of scenes measures once. Selection costs ~130 ms on MPS, ~380 ms on CPU.

Verification

• 24 new tests covering tie-breaking, emulated-dtype screening, unsupported dtypes, caching, probe sizing and passthrough
• 202 tests pass overall, ruff + mypy clean
• Deterministic across 6 fresh processes on MPS (bfloat16) and CPU (float32)

🤖 Generated with Claude Code