CSC format

This notebook introduces CSCArray, the column-compressed companion to CSRArray. CSC stores columns contiguously and is especially useful when column access or transpose products are the natural operation.

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

data = mx.array(np.array([2.0, -1.0, 4.0, 5.0], dtype=np.float32))
row = mx.array(np.array([0, 2, 1, 3], dtype=np.int32))
col = mx.array(np.array([0, 0, 2, 2], dtype=np.int32))

coo = ms.coo_array((data, (row, col)), shape=(4, 3))
csc = coo.tocsc(canonical=True)
csc

CSCArray(shape=(4, 3), nnz=4, dtype=mlx.core.float32, index_dtype=mlx.core.int32, sorted_indices=True, has_canonical_format=True)

CSC uses indices for row ids and indptr for column boundaries. The transpose of a CSC matrix is naturally represented as CSR, so csc.T is a zero-copy CSRArray view of the transposed structure.

print("data   ", np.array(csc.data))
print("indices", np.array(csc.indices))
print("indptr ", np.array(csc.indptr))
print("dense\n", np.array(csc.todense()))
print("transpose type", type(csc.T).__name__, csc.T.shape)

data    [ 2. -1.  4.  5.]
indices [0 2 1 3]
indptr  [0 2 2 4]
dense
 [[ 2.  0.  0.]
 [ 0.  0.  4.]
 [-1.  0.  0.]
 [ 0.  0.  5.]]
transpose type CSRArray (3, 4)

CSR and CSC conversion is structural. CSRArray.tocsc() dispatches the native csr_tocsc count/prefix/fill primitive; CSCArray.tocsr() dispatches the native csc_tocsr inverse. Both preserve values and support float32, float16, bfloat16, and complex64 with int32 or int64 indices.

csr = csc.tocsr()
roundtrip = csr.tocsc()

np.testing.assert_allclose(np.array(csr.todense()), np.array(csc.todense()))
np.testing.assert_allclose(np.array(roundtrip.todense()), np.array(csc.todense()))
csr, roundtrip

(CSRArray(shape=(4, 3), nnz=4, dtype=mlx.core.float32, index_dtype=mlx.core.int32, sorted_indices=False, has_canonical_format=False),
 CSCArray(shape=(4, 3), nnz=4, dtype=mlx.core.float32, index_dtype=mlx.core.int32, sorted_indices=False, has_canonical_format=False))

csc_matvec computes A @ x with a native column scatter-add kernel. csc_matvec_transpose computes A.T @ x with native segmented column reductions, which is the product CSC is naturally best at.

x = mx.array(np.array([1.0, -2.0, 0.5], dtype=np.float32))
xt = mx.array(np.array([1.0, 0.0, -1.0, 2.0], dtype=np.float32))

y = csc @ x
yt = ms.csc_matvec_transpose(csc, xt)

dense = csc.todense()
np.testing.assert_allclose(np.array(y), np.array(dense @ x))
np.testing.assert_allclose(np.array(yt), np.array(mx.transpose(dense) @ xt))
np.array(y), np.array(yt)

(array([ 2. ,  2. , -1. ,  2.5], dtype=float32),
 array([ 3.,  0., 10.], dtype=float32))

SciPy CSC matrices can be imported directly with format="csc". Canonicalization follows SciPy semantics: row indices are sorted within each column and duplicate row entries in a column are summed.

import scipy.sparse as sp

sp_csc = sp.csc_matrix(
    (
        np.array([1.0, 2.0, -3.0], dtype=np.float32),
        np.array([0, 0, 2], dtype=np.int32),
        np.array([0, 2, 3], dtype=np.int32),
    ),
    shape=(3, 2),
)
A = ms.from_scipy(sp_csc, format="csc")

assert isinstance(A, ms.CSCArray)
np.testing.assert_allclose(np.array(A.todense()), sp_csc.toarray())
A

CSCArray(shape=(3, 2), nnz=2, dtype=mlx.core.float32, index_dtype=mlx.core.int32, sorted_indices=True, has_canonical_format=True)