Add sparse matrix interface

[majosm]

Adds a generic interface for operations involving sparse CSR matrices.

cc @lukeolson

[majosm]

@inducer I added a loopy kernel fallback for eager (code), and it's working when the array being multiplied is 1D, but I'm running into some trouble with array strides for higher dimensions. I set up the input array like this:

    u1 = actx.np.where(actx.np.less_equal(x, (a+b)/2), 0*x + 1, 0*x)
    u2 = actx.np.where(actx.np.greater_equal(x, (a+b)/2), 0*x + 1, 0*x)
    # u = actx.freeze_thaw(u1)  # 1D
    u = actx.freeze_thaw(actx.np.stack([u1, u2]).T)  # 2D

The strides are no longer C-like after transposing. Lazy doesn't seem to care about this, but eager doesn't like it; it makes my 1D diffusion solver blow up instantly. If I add order="F" to the GlobalArg(...) calls for "array" and "out", then it works fine. Is there a way I can fix this in general?

Edit: I guess the strides of u are only non-C-like in the eager case.

[majosm]

FWIW, using a dense matrix + actx.einsum instead shows the same behavior.

[majosm]

Upon closer inspection, it looks like the kernels may actually be OK. The problem occurs later in the test driver, when I do this:

        u = actx.freeze_thaw(u + dt*compiled_rhs(t, u))

(here u has order='F' and dt*compiled_rhs(t, u) has order='C' in the eager case). They don't add together correctly in eager, which seems to be due to inducer/pyopencl#681. So I guess the solution is just "don't mix orders".

[majosm]

I added some limited support for applying sparse matrices to array containers. Currently it's restricted to the cases of mat @ container. I thought about also allowing matrices to live inside containers (so one could do container_of_mat @ array or container_of_mat @ container), but decided against it for now because it has some complications:

1. Matrices would need to be arrays. I tried doing this in pytato, and it's not too horrible; I'm not sure if it would work for numpy though (since scipy.sparse.csr_matrix is not a numpy array).
2. @ operations might be ambiguous in some cases (e.g., object arrays define @).

@inducer I think this is ready for a first look. I still have some FIXMEs scattered throughout the changes for things I could use some input on.

[Copilot]

Pull request overview

This PR introduces a new sparse matrix API to arraycontext, centered on a generic SparseMatrix base type and a CSRMatrix implementation, along with actx.make_csr_matrix and actx.sparse_matmul/__matmul__ support across selected backends.

Changes:

• Add SparseMatrix/CSRMatrix types and CSR matmul support to the core ArrayContext interface.
• Implement sparse matmul for NumPy (via SciPy) and Pytato (via pytato.make_csr_matrix); explicitly mark JAX variants as not supported.
• Add CSR matmul tests and introduce a sparse optional dependency group (SciPy), plus CI typing support (scipy-stubs).

Reviewed changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated 6 comments.

Show a summary per file

File	Description
arraycontext/context.py	Adds SparseMatrix/CSRMatrix, CSR construction, and a Loopy-based CSR matmul implementation in the base ArrayContext.
arraycontext/impl/numpy/__init__.py	Implements sparse_matmul for CSRMatrix using scipy.sparse.csr_matrix.
arraycontext/impl/pytato/__init__.py	Implements sparse_matmul for CSRMatrix using Pytato CSR support; adds overrides for JAX-based Pytato context to raise NotImplementedError.
arraycontext/impl/jax/__init__.py	Adds explicit NotImplementedError for sparse matrix creation/matmul.
test/test_arraycontext.py	Adds end-to-end CSR matmul tests over arrays, object arrays, and array containers.
pyproject.toml	Adds sparse = ["scipy"] optional dependency group.
.github/workflows/ci.yml	Adds scipy-stubs to the type-checking install list.
arraycontext/__init__.py	Exports SparseMatrix and CSRMatrix at the package top level.
.basedpyright/baseline.json	Updates type-checking baseline.

Comments suppressed due to low confidence (1)

pyproject.toml:58

• The new sparse optional dependency adds scipy, but the test extra does not include it even though tests now require sparse matmul to work for NumpyArrayContext. This will make default test installs (.[test]) fail unless SciPy happens to be present. Consider adding scipy (or arraycontext[sparse]) to the test extra, or making the tests/implementation gracefully skip when SciPy is not installed.

sparse = [
    "scipy",
]
test = [
    "basedpyright",
    "pytest",
    "ruff",
    # for type checking
    "pymbolic",
]

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

You can also share your feedback on Copilot code review. Take the survey.

[inducer]

These look like reasonable concerns except as noted. @majosm, could you take a look?

[inducer]

Thx!