Skip to content

Device

igl.get_device()

Return the best available torch device.

Preference order: Apple mps, then cuda, then cpu. Each device type has its own execution branch (see :func:select_backend): the CPU branch is the bit-exact reference, the MPS and CUDA branches keep every tensor on the device and solve the readout with an on-device Cholesky.

Returns:

Type Description
device

The selected :class:torch.device.

Execution branches

Training selects one backend per fit from the module's device. The backend owns everything that differs between devices: the readout solve, the Green kernel path, the optimizer construction, snapshots for early stopping, the matmul precision context, and the single host transfer per epoch. See Device branches for the table.

igl.device.select_backend(device, *, tf32=True, cpu_threads=None)

The backend for device: CpuBackend, MpsBackend or CudaBackend.

Parameters:

Name Type Description Default
device device | str

The device the module lives on.

required
tf32 bool

Whether the CUDA backend enables TF32 matmuls during a fit.

True
cpu_threads int | None

Intra-op thread cap for the CPU backend during a fit (None leaves torch's setting; results then stay bit-identical).

None

Returns:

Name Type Description
A Backend

igl.device.Backend

Bases: Protocol

The per-device branch of the training loop.

Attributes:

Name Type Description
name BackendName

The device type this backend serves.

kernel_path KernelPath

Which formulation the Green kernel uses on this device.

igl.device.CpuBackend

The reference branch: bit-identical to the pre-backend code.

threads caps torch's intra-op thread count for the duration of a fit (:meth:precision). Off by default: fewer threads are faster on small problems (an M4 Max runs the medium benchmark batch 13% faster with 6 threads than with its default 12), but the thread count changes the order of BLAS reductions, so results are then no longer bit-identical to the default's.

igl.device.MpsBackend

Bases: _DeviceBackend

Apple MPS: on-device Cholesky solve, float32 accumulation (MPS has no float64).

The readout stays on the device on purpose: factoring the R × R system in float64 on the CPU is 3× faster in isolation, but the copy drains the asynchronous Metal queue every batch and the fit gets 10–20% slower (measured, benchmarks/device/REPORT.md); the hybrid is kept for the one-off public solve only.

igl.device.CudaBackend

Bases: _DeviceBackend