spacr.accelerator

Select the torch compute device and report its vendor name and limitations.

About forty-five spaCR call sites ask torch.cuda.is_available(), and most of them use it to decide whether any GPU is present. On non-NVIDIA hardware those are two different questions, and treating them as one reports a usable accelerator as “no GPU”:

• ROCm answers torch.cuda.is_available() True and wants device="cuda". Code that then prints “CUDA” names the wrong vendor; code that prints “no NVIDIA GPU” hides a usable AMD device.

• Apple’s Metal answers it False and wants device="mps". MPS is Metal, not an Apple-Silicon feature. On the Intel iMac this was written on it drives an AMD Radeon Pro 5300 – measured 16x on conv2d inference and 2.7x on a training step against the same machine’s CPU. ROCm has no macOS build, so Metal is the only route to an AMD card on a Mac.

This module therefore reports the vendor and the GPU decision separately, and also records what the selected device cannot do. Moving a tensor to MPS is not sufficient, because three operations fail rather than slow down (each was measured):

• float64 raises TypeError. Computations that need double precision must stay on the CPU.

• autocast(device_type="mps") raises RuntimeError. Mixed precision must be disabled on this device, not redirected to it.

• Some operators are missing (aten::linalg_qr.out among them). PYTORCH_ENABLE_MPS_FALLBACK=1 makes those run on the CPU instead of raising, so this module sets it when it selects MPS.

Detected and usable hardware are reported separately. A setup screen that lists an accelerator spaCR never dispatches to leads users to attribute CPU run times to that hardware. Neural engines are the main case: the Apple Neural Engine, Intel AI Boost and Qualcomm Hexagon have no portable torch device and are never selected; neural_engines() currently reports only the Apple Neural Engine, on arm64 macOS.

Device resolution does not raise. A half-installed ROCm, a broken driver, or a torch build without a backend it advertises all resolve to the CPU with an explanatory note, so an unusual accelerator cannot cause a traceback at import time.

Classes

Accelerator

Resolved compute device, its display names and its capabilities.

Functions

autocast_device_type(→ Optional[str])	The string for torch.autocast(device_type=...), or None.
capabilities(→ Tuple[Tuple[str, bool, str], ...])	(task, accelerated, detail) rows stating which tasks are accelerated.
cellpose_gpu(→ bool)	Return the value to pass Cellpose as gpu=.
cellpose_kwargs(→ dict)	Consistent device and weight precision for CellposeModel inference.
describe(→ str)	One line for a log or a console: what was found and whether it runs.
device_string(→ str)	The resolved device as a string, e.g. "cuda:0", "mps".
empty_cache(→ str)	Release the backend's cached device memory and return the call made
inspect_torch(→ Accelerator)	Resolve the accelerator for a given torch module, bypassing the cache.
is_cuda(→ bool)	Whether the resolved accelerator is NVIDIA CUDA (not ROCm).
is_gpu(→ bool)	Whether a usable accelerator of any vendor was selected.
neural_engines(→ Tuple[str, ...])	Names of detected neural-inference hardware that spaCR does not use.
resolve(→ Accelerator)	The accelerator spaCR will use. Cached; never raises.
supports_autocast(→ bool)	Whether a GPU is selected and torch.autocast accepts its device type.
supports_bfloat16(→ bool)	Whether torch.bfloat16 tensors may be put on the device.
supports_float64(→ bool)	Whether double precision may be sent to the device.
torch_device()	torch.device for the resolved accelerator.

Module Contents

class spacr.accelerator.Accelerator

Resolved compute device, its display names and its capabilities.

Parameters:

• kind – one of KINDS. rocm is reported distinctly from cuda even though both dispatch to device="cuda", because the only thing a user can act on is the true vendor.

• device – the string to hand torch.device.

• label – human text for the setup slide and the doctor.

• name – undecorated hardware marketing name shown to users, without the backend or version suffix carried by label; empty when no device name was discovered.

• detected – the hardware is present.

• usable – spaCR dispatches work to it. Never true unless detected; false for accelerators with no torch device.

• note – why it is not usable, when it is not. Empty otherwise.

• float64 – the device accepts double precision.

• autocast – torch.autocast accepts this device type.

• fallback – missing operators silently run on the CPU instead of raising, because this module set the backend’s fallback flag.

• bfloat16 – the device accepts torch.bfloat16. Probed rather than assumed from the backend name, because it is the one capability that moves with the torch version: Metal gained it after 2.2, and 2.2.2 is the last x86-64 macOS wheel, so the same backend answers differently on an Intel Mac and an Apple Silicon one. Cellpose’s cpsam loads its weights in bfloat16 by default, so a wrong answer here is a TypeError at model construction.

property is_cuda: bool

Whether this accelerator is NVIDIA CUDA (not ROCm).

Narrower than is_gpu: anything reading nvidia-smi, a CUDA version, or NVIDIA-specific memory interfaces wants this one and would be wrong on ROCm.

property is_gpu: bool

Whether a usable non-CPU device was selected.

This is the check most call sites intend when they write torch.cuda.is_available().

property torch_device

A torch.device for device; requires torch to be installed.

spacr.accelerator.autocast_device_type() → str | None

The string for torch.autocast(device_type=...), or None.

None means “do not use mixed precision here”, which is a different answer from “use it on the CPU” and has to stay distinguishable.

spacr.accelerator.capabilities(found: Accelerator | None = None) → Tuple[Tuple[str, bool, str], ...]

(task, accelerated, detail) rows stating which tasks are accelerated.

The setup screen reports acceleration per task rather than a single “compatible GPU” flag, because backends differ by task: on Metal, segmentation and the classifier are accelerated while the cuML reductions run on the CPU.

Ordered by how much the acceleration is worth: segmentation on the CPU took 444 s for one 256x256 image on the machine this was written on, and 3.2 s on its Radeon.

Parameters:

found – optional existing accelerator snapshot. Providing it avoids another resolution and any dtype-probe allocations, as required by the doctor’s metadata-only check.

Returns:

task, accelerated flag and explanation for each capability.

spacr.accelerator.cellpose_gpu() → bool

Return the value to pass Cellpose as gpu=.

Cellpose resolves the device itself and supports MPS – assign_device(gpu=True) returns mps on a Metal machine – but it uses a GPU only when called with gpu=True. Passing torch.cuda.is_available() here pins a Mac to the CPU when no explicit device is supplied. An explicit Cellpose device takes precedence.

spacr.accelerator.cellpose_kwargs() → dict

Consistent device and weight precision for CellposeModel inference.

The three arguments must agree, so they are produced together here rather than written out at six call sites:

• gpu – controls Cellpose’s selection when no explicit device is supplied, including callers that deliberately drop device below.

• device – from the one resolver, so cellpose and spaCR cannot disagree about the same machine; an explicit device takes precedence.

• use_bfloat16 – cpsam loads its weights in bfloat16 by default and Metal on torch 2.2 has no bfloat16, so the default is a TypeError: BFloat16 is not supported on MPS at construction. Measured on the reporting iMac; float32 weights work there at the cost of additional VRAM. CPU inference also uses float32: bfloat16 support does not imply native arithmetic, and emulation can be substantially slower. This matches Make Masks; float32 arithmetic need not produce identical predictions to bfloat16.

Callers that pass device=None on purpose – letting cellpose resolve it – should take gpu and use_bfloat16 from here and drop device.

spacr.accelerator.describe() → str

One line for a log or a console: what was found and whether it runs.

spacr.accelerator.device_string() → str

The resolved device as a string, e.g. "cuda:0", "mps".

spacr.accelerator.empty_cache(torch_module=None) → str

Release the backend’s cached device memory and return the call made ("" when none was made). Never raises.

Parameters:

torch_module – resolve the backend from this torch module instead of the cached result for the machine. spacr.qt.resource_cleanup holds its own handle – it deliberately does not import torch just to free memory – and a cached global would send the call to the wrong backend when that handle is a stand-in.

spacr.accelerator.inspect_torch(torch, *, device_names=True, include_cuda=True) → Accelerator

Resolve the accelerator for a given torch module, bypassing the cache.

For callers that already hold a torch handle and need the result for that handle – spacr.doctor.check_gpu() imports torch through its own indirection so the diagnosis can be tested against a stand-in, and the cached result for the real machine would ignore that stand-in. Uses the same probes in the same order as resolve().

Parameters:

• torch – torch module whose availability metadata should be inspected.

• device_names – False avoids CUDA/XPU property queries that can initialize their runtime. Availability queries still run.

• include_cuda – False skips CUDA/ROCm entirely, for a process whose environment deliberately hides those devices. Other backends remain eligible. This function does not allocate dtype-probe tensors.

Returns:

accelerator described by the supplied torch module.

spacr.accelerator.is_cuda() → bool

Whether the resolved accelerator is NVIDIA CUDA (not ROCm).

For the places that need CUDA specifically – memory interrogation, nvidia-smi advice, CUDA version reporting.

spacr.accelerator.is_gpu() → bool

Whether a usable accelerator of any vendor was selected.

What a call site means when it writes torch.cuda.is_available() to decide whether to use a GPU at all.

spacr.accelerator.neural_engines() → Tuple[str, ...]

Names of detected neural-inference hardware that spaCR does not use.

Reported so the setup slide can say “found, not used” for an idle Neural Engine. There is no portable torch device for any of these – the ANE is reachable only through CoreML, Intel’s AI Boost only through OpenVINO – so none of them is offered as a compute device.

spacr.accelerator.resolve(refresh: bool = False) → Accelerator

The accelerator spaCR will use. Cached; never raises.

Parameters:

refresh – probe again instead of answering from the cache. For tests, which fake the probes.

spacr.accelerator.supports_autocast() → bool

Whether a GPU is selected and torch.autocast accepts its device type.

spacr.accelerator.supports_bfloat16() → bool

Whether torch.bfloat16 tensors may be put on the device.

spacr.accelerator.supports_float64() → bool

Whether double precision may be sent to the device.

spacr.accelerator.torch_device()

torch.device for the resolved accelerator.

The direct replacement for torch.device("cuda:0" if torch.cuda.is_available() else "cpu").

Nested helpers

_keep_cellpose_flows_off_metal.on_the_cpu(masks, flows, threshold=0.4, device=None)

Run Cellpose flow-error filtering on the CPU for Metal.

Parameters:

• masks – label masks forwarded unchanged to Cellpose.

• flows – flow field forwarded unchanged to Cellpose.

• threshold – flow-error threshold forwarded to Cellpose.

• device – ignored compatibility argument; the wrapped call always receives torch.device("cpu").

Returns:

the original remove_bad_flow_masks return value.

spacr/accelerator.py:734

_measure_dtypes.accepts(name: str) → bool | None

Probe whether the resolved device accepts one named torch dtype.

Parameters:

name – torch dtype attribute name, such as "float64" or "bfloat16".

Returns:

True when a two-element tensor allocates on the device, False for the errors used by unsupported dtypes, or None when torch lacks the dtype or the probe fails unexpectedly.

spacr/accelerator.py:464