Constructors#

These functions create sparse array instances from raw buffers or structured definitions. All array inputs accept MLX arrays, NumPy arrays, or any sequence convertible via mx.array().

csr_array#

mlx_sparse.csr_array(arg, shape, *, validate='metadata', sorted_indices=False, canonical=None)[source]#

Construct a CSRArray from explicit CSR buffers.

Accepts either a (data, indices, indptr) triple or an existing CSRArray. All array inputs are converted to mlx.core.array if they are not already.

Parameters:
  • arg

    A length-3 iterable (data, indices, indptr) where

    • data: non-zero values, shape (nnz,), dtype float32 | float16 | bfloat16 | complex64.

    • indices: column indices, shape (nnz,), dtype int32 | int64.

    • indptr: row pointers, shape (n_rows + 1,), same integer dtype as indices.

    Alternatively, an existing CSRArray instance (returned unchanged if shape matches).

  • shape – Matrix dimensions as a length-2 sequence (n_rows, n_cols).

  • validate (bool | Literal['metadata', 'full']) –

    Validation level, one of:

    • "metadata" (default): checks ranks, lengths, and dtypes without reading array values. Safe to call on device arrays.

    • "full" / True: also verifies indptr monotonicity and column index bounds. May synchronize to host.

    • False / "none": skips all checks.

  • sorted_indices (bool) – Set to True to assert that column indices within each row are already sorted ascending. Default False.

  • canonical (bool | None) – Set to True to assert canonical format (sorted indices, no duplicate columns). Implies sorted_indices=True. Default None (not asserted).

Returns:

A CSRArray with the given buffers and shape.

Raises:
  • TypeError – If arg is not a 3-tuple or a CSRArray, or if dtype constraints are violated.

  • ValueError – If shape or length constraints are violated.

Return type:

CSRArray

Example:

import mlx.core as mx
import mlx_sparse as ms

data = mx.array([1.0, 2.0, 3.0], dtype=mx.float32)
indices = mx.array([0, 1, 0], dtype=mx.int32)
indptr = mx.array([0, 2, 3], dtype=mx.int32)

A = ms.csr_array((data, indices, indptr), shape=(2, 2))
# Full validation from host arrays:
A_checked = ms.csr_array(
    (data, indices, indptr), shape=(2, 2), validate="full"
)

coo_array#

mlx_sparse.coo_array(arg, shape, *, validate='metadata', canonical=None)[source]#

Construct a COOArray from coordinate arrays.

Accepts either a (data, (row, col)) pair or an existing COOArray. All array inputs are converted to mlx.core.array if they are not already.

Parameters:
  • arg

    A (data, (row, col)) pair where

    • data: non-zero values, shape (nnz,), dtype float32 | float16 | bfloat16 | complex64.

    • row: row coordinates, shape (nnz,), dtype int32 | int64.

    • col: column coordinates, shape (nnz,), same integer dtype as row.

    Alternatively, an existing COOArray (returned unchanged if shape matches).

  • shape – Matrix dimensions as a length-2 sequence (n_rows, n_cols).

  • validate (bool | Literal['metadata', 'full']) –

    Validation level, one of:

    • "metadata" (default): checks ranks, lengths, and dtypes.

    • "full" / True: also verifies coordinate bounds. May synchronize to host.

    • False / "none": skips all checks.

  • canonical (bool | None) – Set to True to assert the coordinates are sorted and duplicate-free. Default None (not asserted).

Returns:

A COOArray with the given buffers and shape.

Raises:
  • TypeError – If arg cannot be unpacked as (data, (row, col)), or if dtype constraints are violated.

  • ValueError – If shape or length constraints are violated.

Return type:

COOArray

Example:

import mlx.core as mx
import mlx_sparse as ms

data = mx.array([1.0, 2.0, 3.0, 2.0], dtype=mx.float32)
row = mx.array([0, 0, 1, 0], dtype=mx.int32)
col = mx.array([0, 1, 0, 0], dtype=mx.int32)

# Two entries at (0, 0). Summed when converting to CSR.
coo = ms.coo_array((data, (row, col)), shape=(2, 2))
csr = coo.tocsr(canonical=True)  # sums duplicates

eye#

mlx_sparse.eye(n, m=None, *, k=0, dtype=mlx.core.float32, index_dtype=mlx.core.int32)[source]#

Return a sparse identity-like CSR matrix with ones on a specified diagonal.

Produces the same result as numpy.eye() with k=k, but returns a CSRArray instead of a dense array. The matrix has at most min(n, m) stored values. Rows (or columns) that the diagonal does not pass through are empty rows in the CSR representation.

Parameters:
  • n (int) – Number of rows.

  • m (int | None) – Number of columns. Defaults to n, producing a square matrix.

  • k (int) – Diagonal offset. 0 selects the main diagonal. Positive values shift the diagonal above the main diagonal (superdiagonal). Negative values shift it below (subdiagonal).

  • dtype – Value dtype for the stored ones. Must be one of mx.float32, mx.float16, mx.bfloat16, or mx.complex64. Defaults to mx.float32.

  • index_dtype – Integer dtype for indices and indptr. Must be mx.int32 or mx.int64. Defaults to mx.int32.

Returns:

A canonical CSRArray with has_canonical_format=True and sorted_indices=True.

Raises:

TypeError – If dtype or index_dtype is not a supported value.

Return type:

CSRArray

Example:

import mlx_sparse as ms
import mlx.core as mx

# 4x4 identity matrix
I = ms.eye(4)
mx.eval(I.data)
# CSRArray(shape=(4, 4), nnz=4, ...)

# 3x5 matrix with ones on the first superdiagonal
A = ms.eye(3, 5, k=1)
# Non-zeros at (0,1), (1,2), (2,3)

identity#

mlx_sparse.identity(n, dtype=None, format=None, *, index_dtype=mlx.core.int32)[source]#

Return a sparse square identity matrix.

identity(n) is a SciPy-compatible square alias for mlx_sparse.eye(). The default output format is CSR, matching eye(n). Pass format="coo" or format="csc" to request another supported sparse format. Unsupported SciPy formats such as "dia" and "bsr" are rejected explicitly.

Parameters:
  • n (int) – Number of rows and columns. Must be non-negative.

  • dtype – Stored value dtype. None defaults to mx.float32.

  • format – Output format, one of None, "csr", "coo", or "csc". None returns CSR.

  • index_dtype – Integer dtype for sparse indices, mx.int32 or mx.int64.

Returns:

A sparse square identity matrix in the requested format.

diags#

mlx_sparse.diags(diagonals, offsets=0, *, shape=None, dtype=None, index_dtype=mlx.core.int32)[source]#

Construct a CSR matrix from one or more diagonals.

Mirrors the behaviour of scipy.sparse.diags() but returns a CSRArray. Each diagonal is placed at the position specified by the corresponding offset. Diagonals are assembled into a COO triple and sorted before the CSR row-pointer array is built, so the result is always in canonical form. When diagonal values are MLX arrays, the fixed diagonal topology is assembled without converting those values to NumPy, so mx.jvp and mx.vjp propagate sparse-value tangents and cotangents.

Parameters:
  • diagonals

    The diagonal values. Accepted forms:

    • A single 1-D array-like (or scalar) placed at offsets.

    • A 2-D array whose rows are individual diagonals.

    • A list of 1-D array-likes, one per entry in offsets.

    Each diagonal’s length must not exceed the number of elements that the diagonal at the corresponding offset can hold given shape.

  • offsets – Diagonal offset(s). 0 is the main diagonal. Positive integers are superdiagonals. Negative integers are subdiagonals. When diagonals is a list, offsets must be a matching list of integers. Repeated offsets are not allowed.

  • shape (Sequence[int] | None) – Output matrix shape as (n_rows, n_cols). When omitted, the minimum square shape that fits all diagonals is inferred automatically.

  • dtype – Value dtype. When None (default), the dtype is inferred from the diagonal arrays: complex64 if any diagonal is complex, float16 if any diagonal has dtype float16, otherwise float32.

  • index_dtype – Integer dtype for indices and indptr. Must be mx.int32 or mx.int64. Defaults to mx.int32.

Returns:

A canonical CSRArray with has_canonical_format=True and sorted_indices=True.

Raises:
  • TypeError – If dtype or index_dtype is not supported.

  • ValueError – If the number of diagonals and offsets differ, if offsets are repeated, or if a diagonal is longer than its allocated space.

Return type:

CSRArray

Example:

import numpy as np
import mlx_sparse as ms
import mlx.core as mx

# Tridiagonal matrix: main diagonal 2, off-diagonals -1
A = ms.diags(
    [np.full(4, -1.0), np.full(5, 2.0), np.full(4, -1.0)],
    offsets=[-1, 0, 1],
)
# 5x5, nnz=13

# Single diagonal at offset 2
B = ms.diags([1.0, 2.0, 3.0], offsets=2, shape=(5, 5))

block_array#

mlx_sparse.block_array(blocks, *, format=None, dtype=None)[source]#

Build a sparse array from a rectangular grid of sparse or dense blocks.

block_array mirrors the SciPy structural constructor while returning mlx-sparse arrays. Each non-None entry must be a COO, CSR, CSC, or dense rank-2 array. None entries represent all-zero blocks whose shape is inferred from the other blocks in the same block row and block column; all-None rows or columns are assigned size zero.

The block grid is validated before assembly: every block row must have a consistent height, every block column must have a consistent width, the grid must be rectangular, and only "coo", "csr", and "csc" output formats are supported. Dense blocks are converted with the native fromdense() path. Sparse blocks are converted to COO through native format conversion, then one native coordinate-offset primitive copies values and offsets coordinates. No stored entries are iterated in Python.

format=None defaults to COO because block assembly is a construction operation. CSR and CSC requests canonicalize through the existing native compressed conversion path, summing duplicate coordinates.

Parameters:
  • blocks – Rectangular 2-D grid of sparse arrays, dense rank-2 arrays, or None entries.

  • format – Output format, one of None, "coo", "csr", or "csc".

  • dtype – Optional output value dtype. When omitted, value dtypes are promoted across non-None blocks under mlx-sparse’s sparse value dtype policy.

Returns:

A COOArray, CSRArray, or CSCArray.

bmat#

mlx_sparse.bmat(blocks, format=None, dtype=None)[source]#

Compatibility alias for block_array().

Unlike SciPy’s historical matrix-returning bmat behavior, mlx-sparse always returns sparse array containers (COO, CSR, or CSC). See block_array() for validation, dtype promotion, and native assembly details.

block_diag#

mlx_sparse.block_diag(mats, format=None, dtype=None)[source]#

Build a block diagonal sparse array from a sequence of matrices.

Each input must be a sparse COO/CSR/CSC array or a dense rank-2 array. Blocks are placed on the main block diagonal using native coordinate-offset assembly; off-diagonal zero regions are implicit and never materialized. None is not accepted because it has no shape in this constructor.

format=None defaults to COO, matching SciPy’s construction-oriented default. CSR and CSC requests canonicalize through native conversion.

Parameters:
  • mats – Non-empty sequence of sparse or dense rank-2 matrices.

  • format – Output format, one of None, "coo", "csr", or "csc".

  • dtype – Optional output value dtype.

Returns:

A sparse block diagonal array.

vstack#

mlx_sparse.vstack(blocks, format=None, dtype=None)[source]#

Stack sparse or dense rank-2 blocks vertically.

All blocks must have the same number of columns. The implementation uses the same native coordinate-offset assembly as block_array() with a single block column. None entries are rejected because their row height cannot be inferred in a one-dimensional stack.

Parameters:
  • blocks – Non-empty sequence of sparse or dense rank-2 matrices.

  • format – Output format, one of None, "coo", "csr", or "csc". None defaults to COO.

  • dtype – Optional output value dtype.

Returns:

A sparse array containing the vertical stack.

hstack#

mlx_sparse.hstack(blocks, format=None, dtype=None)[source]#

Stack sparse or dense rank-2 blocks horizontally.

All blocks must have the same number of rows. The implementation uses native coordinate-offset assembly with a single block row. None entries are rejected because their column width cannot be inferred in a one-dimensional stack.

Parameters:
  • blocks – Non-empty sequence of sparse or dense rank-2 matrices.

  • format – Output format, one of None, "coo", "csr", or "csc". None defaults to COO.

  • dtype – Optional output value dtype.

Returns:

A sparse array containing the horizontal stack.

tril#

mlx_sparse.tril(A, k=0, format=None)[source]#

Return the lower triangular portion of a sparse or dense matrix.

Elements with column - row <= k are retained. k=0 keeps the main diagonal, positive k includes superdiagonals, and negative k moves the cutoff below the main diagonal. COO, CSR, CSC, and dense rank-2 inputs are accepted. Dense inputs are first extracted with native fromdense(); sparse inputs are filtered with native staged count/fill kernels for their storage format.

format=None defaults to COO, matching SciPy’s triangular extraction default. CSR and CSC requests are returned in the requested sparse format. Triangular extraction compacts sparse structure based on coordinates and is intentionally forward-only under MLX autodiff in v0.0.6b0.

triu#

mlx_sparse.triu(A, k=0, format=None)[source]#

Return the upper triangular portion of a sparse or dense matrix.

Elements with column - row >= k are retained. See tril() for input handling, native compaction details, format semantics, and autodiff limitations.

fromdense#

mlx_sparse.fromdense(array, *, threshold=0.0, dtype=None, index_dtype=mlx.core.int32)[source]#

Construct a canonical CSR matrix from a rank-2 dense MLX array.

Identifies the non-zero (or above-threshold) entries of a dense matrix and packages them into a CSRArray. The native path stages this as count, allocate, then fill work so Metal builds can perform the dense scan and CSR writes on device while still returning compact buffers. Because the output sparse topology depends on numerical values and the threshold, fromdense is intentionally not differentiable; use fixed- topology constructors when sparse values need gradients.

The value dtype is preserved from the input array. Index dtype defaults to int32 and can be overridden for matrices with more than ~2 billion non-zeros (not typical on Apple Silicon).

Parameters:
  • array – A rank-2 array-like. Converted to mlx.core.array if not already. Dtype must be one of float32, float16, bfloat16, or complex64.

  • threshold (float) – Entries with absolute value less than or equal to threshold are treated as structural zeros and excluded from the output. The default 0.0 keeps every numerically non-zero entry. Must be non-negative.

  • dtype – Optional value dtype to cast to before extracting non-zeros. When None, the input dtype chosen by MLX is preserved.

  • index_dtype – Integer dtype for indices and indptr. Must be mx.int32 or mx.int64. Defaults to mx.int32.

Returns:

A canonical CSRArray with has_canonical_format=True and sorted_indices=True.

Raises:
  • TypeError – If the input dtype is not a supported value dtype.

  • ValueError – If the input is not rank-2, or if threshold is negative.

Return type:

CSRArray

Example:

import mlx.core as mx
import numpy as np
import mlx_sparse as ms

dense = mx.array(np.array([
    [1.0, 0.0, 2.0],
    [0.0, 0.0, 0.0],
    [3.0, 4.0, 0.0],
], dtype=np.float32))

csr = ms.fromdense(dense)
# CSRArray(shape=(3, 3), nnz=4, dtype=float32, ...)

# Drop near-zero entries below 0.1
csr_thresholded = ms.fromdense(dense, threshold=0.5)

from_dense#

mlx_sparse.from_dense(array, *, threshold=0.0, dtype=None, index_dtype=mlx.core.int32)[source]#

Alias for fromdense() with a PEP 8 compatible name.

Parameters:

threshold (float)

Return type:

CSRArray

from_numpy#

mlx_sparse.from_numpy(array, *, threshold=0.0, dtype=None, index_dtype=mlx.core.int32)[source]#

Convert a rank-2 NumPy array to a canonical CSRArray.

Parameters:

threshold (float)

Return type:

CSRArray

from_scipy#

mlx_sparse.from_scipy(matrix, *, format='csr', dtype=None, index_dtype=mlx.core.int32, canonical=True)[source]#

Convert a SciPy sparse matrix or sparse array to mlx-sparse.

Any SciPy sparse format is accepted. format="csr" returns a CSRArray, format="csc" returns a CSCArray, and format="coo" returns a COOArray. The conversion preserves supported float32, float16, and complex64 values. Other real floating dtypes, including SciPy’s default float64, are cast to float32 unless dtype is provided.

Parameters:
  • matrix – A scipy.sparse matrix or array.

  • format (str) – Output sparse format: "csr" (default), "csc", or "coo".

  • dtype – Optional MLX value dtype. Must be one of mx.float32, mx.float16, mx.bfloat16, or mx.complex64.

  • index_dtype – Integer dtype for sparse indices. Must be mx.int32 or mx.int64.

  • canonical (bool) – If True (default), sum duplicates and sort indices before exporting buffers.

Returns:

A CSRArray, CSCArray, or COOArray.

Raises:
  • TypeError – If SciPy is not installed, matrix is not sparse, or a dtype is unsupported.

  • ValueError – If format is not "csr", "csc", or "coo".

asarray#

mlx_sparse.asarray(x, *, threshold=0.0, dtype=None, index_dtype=mlx.core.int32)[source]#

Convert common sparse or dense inputs to a sparse array.

Existing CSRArray and CSCArray instances are returned unchanged unless dtype requests a value cast. COOArray instances are converted with tocsr(canonical=True). SciPy sparse matrices/arrays route through from_scipy(), dense MLX, NumPy, and Python array-likes route through fromdense().

Parameters:
  • x – Existing mlx-sparse array, SciPy sparse array, dense MLX array, NumPy array, or Python rank-2 array-like.

  • threshold (float) – Dense-only structural-zero threshold.

  • dtype – Optional target value dtype.

  • index_dtype – Target index dtype for newly constructed sparse arrays.

Returns:

Existing CSRArray or CSCArray inputs are preserved. Other inputs return a canonical CSRArray.

Return type:

CSRArray | CSCArray

Validation modes#

The validate parameter on csr_array() and coo_array() accepts:

Value

Behaviour

"metadata"

Checks ranks, array lengths, and dtypes. No value reads. Default.

"full" or True

Full metadata checks plus value-level checks (bounds, monotonicity). May call mx.eval to read index values from device.

False or "none"

No checks. Use only when inputs are known valid.

See Validation for a detailed discussion.

Structured constructors#

eye(), identity(), diags(), block_array(), bmat(), block_diag(), vstack(), hstack(), tril(), triu(), fromdense(), from_dense(), from_numpy(), from_scipy(), and asarray() all make common construction paths explicit.

Block assembly and triangular extraction are native-backed and avoid Python loops over stored entries. None entries in block_array() represent implicit zero blocks. Dense blocks are converted with the native fromdense() path before sparse assembly. CSR and CSC format requests use native compressed conversion and canonicalization where duplicate coordinates may need to be summed.

import numpy as np
import mlx.core as mx
import mlx_sparse as ms

# 5x5 identity
I = ms.eye(5)

# 4x4 tridiagonal
T = ms.diags(
    [np.full(3, -1.0), np.full(4, 2.0), np.full(3, -1.0)],
    offsets=[-1, 0, 1],
)

# Convert a dense weight matrix
W = mx.array(np.random.randn(8, 8).astype(np.float32))
W_sparse = ms.fromdense(W, threshold=0.1)

# Convert SciPy sparse or dense NumPy inputs without hand-building buffers
W_from_np = ms.asarray(np.eye(8, dtype=np.float32))

# Assemble a saddle-point-style block matrix
K = ms.block_array([[T, None], [None, ms.identity(4)]], format="csr")

# Extract a triangular sparse view without densifying
L = ms.tril(K, format="csr")