Skip to content

Device Routing#

With CUDA JAX installed, JAX places everything on the GPU — including the small, dispatch-bound programs where a CPU finishes several times faster. hamon's entry points therefore take a device argument:

  • "auto" (default) — with no accelerator visible, placement is untouched. Otherwise the work score (n_chains × free nodes) decides: small workloads run on the CPU, large ones on the accelerator.
  • "cpu" / "gpu" — that platform, raising if it is not visible.
  • a concrete jax.Device — used as-is.
  • None — hamon never touches placement.

The default threshold (4096, the steady-state crossover measured on an RTX 5080) can be overridden with HAMON_DEVICE_THRESHOLD; calibrate your own with python benchmarks/device_crossover.py. HAMON_DEVICE=cpu|gpu|none forces a choice without code changes.

Very short one-shot flows are compile-dominated and can favor the CPU regardless of size — pass device="cpu" for those, or enable the persistent compile cache so repeated runs skip GPU compilation entirely.

hamon.resolve_device(device: DeviceLike = 'auto', *, score: int | None = None, threshold: float | None = None) -> JaxDevice | None #

Resolve a device spec into a concrete jax.Device or None.

None means "leave placement alone" and is what "auto" resolves to when no accelerator is visible, so CPU-only installs behave identically with or without routing. With an accelerator present, "auto" routes to it when score meets the threshold and to the CPU otherwise (or when no score is supplied). Explicit "cpu"/"gpu"/"tpu" requests fail loudly if the device is absent.

hamon.work_score(n_chains: int, n_nodes: int) -> int #

The routing heuristic's work estimate: width of the parallel front.

Rounds are deliberately excluded — the round loop is a single jitted scan, so per-round dispatch does not scale with round count, and compile cost is amortized by the persistent compilation cache.

hamon.enable_persistent_compile_cache(path: str | None = None) -> str | None #

Turn on JAX's persistent compilation cache (idempotent).

XLA compile dominates the cold cost of NRPT and especially the multi-probe autotuning search (each chain count and each n_expl recompiles the round loop). The persistent cache stores compiled executables on disk and reuses them across processes — measured ≈ −72% wall on repeat cold runs — which is what keeps autotuning affordable.

Only enabled on an accelerator backend. On a CPU-only backend XLA's AOT loader logs a "machine-feature mismatch" error (warning of a theoretical SIGILL) for every reloaded executable, and CPU compiles are cheap, so caching there is net-negative; the test conftest makes the same accelerator-only choice. An explicit JAX_COMPILATION_CACHE_DIR overrides this on any backend (the user's deliberate opt-in/out wins).

If JAX_COMPILATION_CACHE_DIR is already set in the environment its directory is respected (including opting out via an empty value). Otherwise, when an accelerator is present, the cache dir is set to path (or ~/.cache/jax by default), matching the GPU default used in the test suite. Whenever caching is active the persistence thresholds are lowered (see :func:_lower_persistence_thresholds).

Parameters:

Name Type Description Default
path str | None

cache directory; defaults to ~/.cache/jax.

 None

Returns:

Type Description
str | None

The active cache directory, or None if caching is disabled (CPU-only
str | None

backend, or an empty env var).

hamon.DeviceLike module-attribute #

Represent a PEP 604 union type

E.g. for int | str