feat: automatic CPU/GPU device routing

[dek3rr]

What

Automatic CPU/GPU device routing: every public entry point (nrpt, nrpt_adaptive, discover_chain_count, ising_sample, sample_states, sample_with_observation, estimate_moments, estimate_kl_grad) now takes device="auto". With no accelerator visible, placement is untouched (byte-for-byte legacy behavior — CPU-only CI sees no change). With a GPU present, the work score n_chains × free_nodes routes small workloads to the CPU and large ones to the GPU.

• hamon/device.py owns all policy: resolve_device, tree_device_put (identity fast path keeps jit caches warm), tracer guard (routing is a no-op inside jit/vmap/grad), HAMON_DEVICE / HAMON_DEVICE_THRESHOLD env overrides, device="cpu"|"gpu"|jax.Device|None escape hatches.
• Orchestrators resolve the device once and pass the concrete device down, so it never flips between tuning phases; the jit-once round loop is preserved (new trace-count regression test for the routed path).
• Test suite pins the default device to CPU (tests/conftest.py); a new gpu marker runs a smoke subset on real hardware (auto-skipped when absent); HAMON_TEST_DEVICE=gpu unpins the suite.
• benchmarks/device_crossover.py measures a machine's crossover and recommends a threshold.

Why

With CUDA jax installed, JAX defaults everything to the GPU — including tiny dispatch-bound programs where the CPU is several times faster. Measured on an RTX 5080 (WSL, jax 0.10.1): the test suite ran 837s on GPU vs 172s CPU-forced, while large sampling is 2–11× faster on GPU. Goal: having GPU jax installed should never make a workload slower than CPU-only jax.

Measured results (RTX 5080)

Workload	Before	After
Full test suite on the GPU machine	837s	159s
Crossover sweep	—	score ≤ 2048 → CPU wins; score ≥ 4096 → GPU wins 2–11×
Auto-routing sanity	—	score 256 → cpu:0, score 32768 → cuda:0

The default threshold (4096) is the measured steady-state crossover. Known simplification: rounds are excluded from the score, so very short one-shot flows remain compile-dominated and can favor CPU — documented escape hatches (device="cpu", JAX_COMPILATION_CACHE_DIR).

Reviewer notes

• Routing re-commits entry arrays to the resolved device; outputs come back committed there. device=None is the contract for "hamon never touches placement".
• tree_device_put maps only jax.Array leaves; leaf-level statics (spec/blocks/nodes) keep object identity and reconstructed container Modules compare equal, so equinox caches hit — proven by test_adaptive_phases_compile_once_with_explicit_cpu.
• Arrays closed over by energy_delta_fn cannot be moved by routing (documented in the nrpt docstring).
• Verified green: Windows CPU-only (296 passed, 3 gpu-skipped), WSL+5080 full suite, WSL -m gpu smoke subset, WSL with JAX_PLATFORMS=cpu.
• The first two commits (float32-under-x64 fix, test compile cache) target main locally but weren't pushed yet; once main is pushed this PR reduces to the two routing commits.

🤖 Generated with Claude Code