Repository navigation
feat(sampler): choose the sampling device with a device argument - #293
Conversation
compile_sampler, compile_detector_sampler and the sampler classes take device="auto" (default), "cpu" or "metal". "auto" samples on an Apple GPU when it can run the circuit and otherwise on the CPU, now without a warning; "cpu" stays on the CPU; "metal" raises with the reason where the GPU cannot run the circuit (no Apple GPU or PyObjC's Metal, forked process, components wider than 11 active qubits, unresolvable rotations, noise beyond the GPU tables, out of GPU memory, kernel compile errors), both at compile time and when a call fails. TSIM_DEVICE=cpu keeps "auto" on the CPU; an explicit device overrides it. The stabilizer-rank strategies take only "auto". Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Codex Review Summary: ✅ CompletedThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
The user skill explains device="auto"/"cpu"/"metal", checking sampler.stats.device since "auto" no longer warns, why GPU calls need 10^6-10^7 shots and batch_size should stay unset (2^20-shot launches that batch_size can only lower), one process per GPU and spawn over fork, and that GPU seeds depend on batch_size and the machine. The benchmarking, profiling and testing skills get the silent fallbacks, the call-size caveat for GPU rates, and device="metal"/"cpu" in Metal tests. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9a6461621e
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| reason = runtime.unavailable() | ||
| if reason is not None: | ||
| if required: | ||
| raise RuntimeError(f"Cannot sample on Metal: {reason}.") |
There was a problem hiding this comment.
When device="metal" is used for a program with no shot words, such as Circuit("M 0").compile_detector_sampler(device="metal"), this availability check raises on a non-Mac before the num_words == 0 fast path is reached. Such a sampler has no work to place on a GPU and is explicitly intended to remain on the host either way, so the result currently varies by machine unnecessarily; check host.layout.num_words before requiring Metal availability.
Useful? React with 👍 / 👎.
| if runtime.forked_after_init(): | ||
| if self._required: | ||
| raise RuntimeError(f"Cannot sample on Metal: {FORKED}.") | ||
| return on_cpu() |
There was a problem hiding this comment.
For an automatic sampler created before a process forks, every call in the child takes this CPU path permanently, but _cpu_only remains false. Consequently sampler.stats.device and the repr continue to report metal even though all samples—including their seed stream and performance characteristics—come from the CPU, which is especially misleading now that this fallback is silent; mark the child sampler CPU-only before returning.
Useful? React with 👍 / 👎.
☂️ Code Coverage
Overall Coverage
New FilesNo new covered files... Modified Files
|
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
…rect Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
A sampler compiled before a fork samples on the CPU in the child, but its stats.device and repr still said "metal"; now that "auto" no longer warns, they must say where it samples. Also correct the docs and skills: a GPU batch is up to 2^20 shots, a GPU call's fixed cost is setting up launches and waiting (Metal buffers are shared, nothing is copied back), TSIM_DEVICE is read at compile time, and the rotation cutoff is about 6e-5 π. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Summary
compile_sampler,compile_detector_sampler,CompiledMeasurementSamplerandCompiledDetectorSamplertake adeviceargument forstrategy="active_statevector":device"auto"(default)"cpu""metal""metal"raises where the GPU cannot run the circuit, at compile time and during a call:RuntimeError: no Apple GPU, PyObjC's Metal not installed, or a process forked after its parent used the GPUValueError: a component wider than 11 active qubits, a rotation branch below float32's resolution (tape.RESOLVABLE), or noise beyond the GPU tablesMemoryError: out of GPU memory, at setup or during a callruntime.CompileError: a Metal kernel that does not compileUnder
"auto", every one of those falls back to the CPU silently;repr(sampler)andsampler.stats.devicesay where it samples. Before this PR,"auto"was the only behaviour and most fallbacks warned.Other details:
TSIM_DEVICE=cpustill keeps"auto"on the CPU, so the benchmark harness and existing scripts are unchanged. An explicitdeviceoverrides it."auto"; they run wherever JAX does."metal"."metal"sampler stays required when it is loaded.deviceis a new keyword argument with an unchanged default, so this is a minor change under semantic versioning."cuda"can be added when the CUDA backend (CUDA backend #219) lands.install.mdx), the user skill andCHANGELOG.mdare updated. The docs also say that with"auto"a seed reproduces samples only on the same kind of machine; usedevice="cpu"to reproduce across machines.Skills
skills/simulating-circuits-with-tsim):devicevalues, and checkingsampler.stats.devicesince"auto"is silent"metal"(long runs that must not drop to the CPU) and"cpu"(reproducibility, comparisons)batch_sizeunset on the GPU: 2^20-shot launches that it can only lowerbatch_sizeand on the kind of machinebenchmarking-tsim: the silent fallbacks (including tiny rotations),TSIM_DEVICEversusdevice=, and why GPU rates from small calls or a smallbatch_sizemeasure overhead.profiling-tsim: the same fallback and overhead notes.testing-tsim-correctness: the seed invariant holds on the CPU only; Metal tests usedevice="metal"against adevice="cpu"reference.Test plan
uv run tsim-dev check: lint and 2776 tests passed on an M4 Pro, including the Metal integration tests.device="cpu", unknown devices, a non-"auto"device withcat5, and"metal"where Metal is unavailable (raises), with"auto"falling back without warnings.TSIM_DEVICE) now check that"auto"is silent and"metal"raises.test/conftest.pykeeps the non-Metal tests on the CPU through the newruntime.unavailable().Not run:
just doc-build, and no benchmark. The sampling paths and kernels are unchanged; only the selection logic and the warnings changed.🤖 Generated with Claude Code