Runtime#

mlx_sparse.runtime exposes enum-backed CPU runtime controls for code that needs reproducible performance settings without spelling raw configuration strings. It is the preferred public interface for thread counts and operation-family parallel gates.

import mlx_sparse as ms

print(ms.runtime.N_THREADS)

ms.runtime.N_THREADS = 8
ms.runtime.SPGEMM_PARALLEL = True
ms.runtime.SPGEMM_THREADS = "inherit"
ms.runtime.SOLVER_PARALLEL = False
ms.runtime.SOLVER_THREADS = 2

with ms.runtime.context(n_threads=1):
    C = A @ B

report_metadata = ms.runtime.info()

Direct controls#

The common runtime controls are module attributes. Reading them returns the effective value that kernels should use. Assigning them validates and updates the underlying package configuration and synchronized MLX_SPARSE_* environment variable.

For native CPU same-format CSR, COO, and CSC sparse-sparse products, SPGEMM_THREADS is a fixed worker count. The implementation partitions independent output rows or columns across that count and does not change the number of workers based on matrix size, density, or estimated work. Very small outputs may assign empty ranges to some workers rather than silently reducing the configured count. Use SPGEMM_THREADS = 1 or SPGEMM_PARALLEL = False to force the serial Gustavson/SPA path.

Attribute	Read value	Accepted writes
`N_THREADS`	Resolved package-wide CPU worker count.	Positive integer or `"auto"`.
`SPGEMM_PARALLEL`	Whether CPU SpGEMM parallelism is enabled.	Boolean-like value.
`SPGEMM_THREADS`	Effective CPU worker count for sparse-sparse products.	Positive integer, `"auto"`, or `"inherit"`.
`SOLVER_PARALLEL`	Whether CPU solver parallelism is enabled.	Boolean-like value.
`SOLVER_THREADS`	Effective CPU worker count for solver routines.	Positive integer, `"auto"`, or `"inherit"`.

Enum keys#

The enum remains available for structured helper calls such as context(RuntimeOption.N_THREADS, 1) and for code that wants stable option identifiers without relying on strings.

class mlx_sparse.runtime.RuntimeOption(*values)[source]#

Bases: str, Enum

Runtime option identifiers accepted by mlx_sparse.runtime.

N_THREADS = 'CPU_THREADS'#

SPGEMM_PARALLEL = 'SPGEMM_PARALLEL'#

SPGEMM_THREADS = 'SPGEMM_THREADS'#

SOLVER_PARALLEL = 'SOLVER_PARALLEL'#

SOLVER_THREADS = 'SOLVER_THREADS'#

Scoped overrides#

Use the context manager when a benchmark or experiment needs temporary runtime settings without permanently changing the process configuration.

mlx_sparse.runtime.context(arg1=None, arg2=<object object>, **kwargs)[source]#

Temporarily patch runtime options.

Accepted forms mirror mlx_sparse.config.patch():

context(ms.runtime.RuntimeOption.N_THREADS, 4): Patch one enum option.
context({ms.runtime.RuntimeOption.N_THREADS: 4}): Patch several enum options.
context(n_threads=4, spgemm_parallel=False): Patch options with readable keyword aliases.

Parameters:

arg1 (RuntimeOption | str | Mapping[Any, Any] | None)
arg2 (Any)
kwargs (Any)

Return type:

Iterator[None]

mlx_sparse.runtime.patch(arg1=None, arg2=<object object>, **kwargs)#

Temporarily patch runtime options.

Accepted forms mirror mlx_sparse.config.patch():

context(ms.runtime.RuntimeOption.N_THREADS, 4): Patch one enum option.
context({ms.runtime.RuntimeOption.N_THREADS: 4}): Patch several enum options.
context(n_threads=4, spgemm_parallel=False): Patch options with readable keyword aliases.

Parameters:

arg1 (RuntimeOption | str | Mapping[Any, Any] | None)
arg2 (Any)
kwargs (Any)

Return type:

Iterator[None]

Thread resolution#

mlx_sparse.runtime.resolve_n_threads()[source]#

Resolve the effective CPU worker count and the source used.

Explicit MLX_SPARSE_CPU_THREADS / ms.runtime.N_THREADS = ... values win first. In "auto" mode, standard thread hints are consulted before scheduler allocations, then process affinity, then hardware concurrency.

Return type:: tuple[int, str]

mlx_sparse.runtime.resolve_spgemm_threads()[source]#

Resolve the CPU worker count and source for sparse-sparse products.

Return type:: tuple[int, str]

mlx_sparse.runtime.resolve_solver_threads()[source]#

Resolve the CPU worker count and source for solver routines.

Return type:: tuple[int, str]

Diagnostics#

mlx_sparse.runtime.info()[source]#

Return structured runtime information for reports and diagnostics.

Return type:: dict[str, Any]