Device selection#

These functions set MLX’s default device and return the selected mlx.core.Device object. They are thin wrappers around mx.set_default_device with an optional availability probe.

use_gpu#

mlx_sparse.use_gpu(index=0, *, require_available=True)[source]#

Set MLX’s default device to GPU (Metal) and return the device object.

On Apple Silicon, this selects the integrated GPU via MLX’s Metal backend. Fixed-shape sparse primitives dispatch native Metal kernels for supported value and index dtypes. Operations with dynamic output structure, such as sparse-sparse products and duplicate summation, may synchronize to host for structural assembly.

Parameters:

  • index (int) – GPU device index. 0 selects the primary GPU. Default 0.

  • require_available (bool) – If True (default), verify that mx.is_available returns True for the selected device and that a trivial array can be evaluated. Raises RuntimeError if either check fails.

Returns:

The mlx.core.Device object that was set as the default.

Raises:

RuntimeError – If require_available=True and the GPU is not available or cannot be probed.

Return type:

mlx.core.Device

Example:

import mlx_sparse as ms

ms.use_gpu()  # pin to GPU
y = A @ x  # dispatches Metal csr_matvec kernel

use_cpu#

mlx_sparse.use_cpu(index=0, *, require_available=True)[source]#

Set MLX’s default device to CPU and return the device object.

Calling this once at the start of a script or test is the recommended way to pin all subsequent MLX and mlx-sparse operations to the CPU. The setting persists for the lifetime of the Python process or until overridden by another use_* call.

Parameters:

  • index (int) – CPU device index. Almost always 0. Multiple CPU devices are not typical in Apple Silicon runtimes. Default 0.

  • require_available (bool) – If True (default), probe the device immediately with a trivial mx.eval and raise RuntimeError if it fails. Set to False to skip the probe (e.g. in environments where eager evaluation is not yet possible).

Returns:

The mlx.core.Device object that was set as the default.

Raises:

RuntimeError – If require_available=True and the CPU device cannot be probed successfully.

Return type:

mlx.core.Device

Example:

import mlx_sparse as ms

ms.use_cpu()  # pin to CPU
y = A @ x  # runs on CPU

use_device#

mlx_sparse.use_device(name, index=0)[source]#

Set MLX’s default device by name string.

A convenience wrapper around use_cpu() and use_gpu() that accepts a plain string device name. Useful when the target device is provided as a command-line argument.

Parameters:

  • name (str) – "cpu" or "gpu" (case-insensitive).

  • index (int) – Device index. Default 0.

Returns:

The mlx.core.Device object that was set as the default.

Raises:

ValueError – If name is not "cpu" or "gpu".

Return type:

mlx.core.Device

Example:

import argparse
import mlx_sparse as ms

parser = argparse.ArgumentParser()
parser.add_argument("--device", default="gpu")
args = parser.parse_args()
ms.use_device(args.device)

Native capabilities#

mlx-sparse exposes enum-backed runtime capability checks for native backend dispatch. Use these when application code needs to branch on whether a backend is actually usable, rather than checking platform strings or package filenames directly.

import mlx_sparse as ms

if ms.capabilities.METAL:
    ms.use_gpu()

if ms.capabilities.ACCELERATE:
    pass  # optional Apple Accelerate sparse direct solvers are available

The public capability names are "extension", "cpu", "metal", "accelerate", "cuda", and "rocm". The corresponding uppercase attributes on capabilities return booleans, for example ms.capabilities.CPU and ms.capabilities.METAL. Status strings are "available", "unavailable", and "not_built".

ms.capabilities.ACCELERATE reports Accelerate-backed sparse solver availability, not just framework presence. Current published macOS wheels (v0.0.4b0 onwards) are built with Accelerate sparse direct-solver support enabled. Linux wheels are CPU-only and report "not_built" for Metal, Accelerate, CUDA, and ROCm. Editable and source macOS builds can opt into the same Accelerate path with MLX_SPARSE_ENABLE_ACCELERATE=ON.

mlx_sparse.capabilities = mlx_sparse.capabilities(extension='unavailable', cpu='unavailable', metal='unavailable', accelerate='unavailable', cuda='unavailable', rocm='unavailable')#

User-facing native capability view.

mlx_sparse.capabilities is intentionally small and string-friendly:

  • capability names: "extension", "cpu", "metal", "accelerate", "cuda", "rocm"

  • statuses: "available", "unavailable", "not_built"

Example:

import mlx_sparse as ms

if ms.capabilities.METAL:
    ms.use_gpu()

if ms.capabilities.status("accelerate") == "not_built":
    ...
mlx_sparse.has_capability(capability)[source]#

Return True if capability is available in this process.

Parameters:

capability (NativeCapability | str)

Return type:

bool

Usage notes#

  • Call one of these functions once, early in your script, to pin execution to a specific device.

  • The selected device persists for the lifetime of the Python process or until another use_* call.

  • The Metal GPU path (use_gpu) requires macOS with Apple Silicon and MLX ≥ 0.31. On unsupported hardware use_gpu raises RuntimeError when require_available=True (the default).

  • Linux builds currently support native CPU execution only. CUDA and ROCm are reserved capability names for future releases.

  • To check sparse native backend availability without setting the default device, prefer ms.capabilities.

See Device selection and execution model for a full discussion of the lazy execution model and which operations run on GPU vs CPU.