`jax.experimental.sparse.CSR` accepts wrong shapes

[johannahaffner]

Description

The following code reproduces the error:

import jax.numpy as jnp
from jax.experimental.sparse import CSR


matrix = jnp.eye(2)
csr = CSR.fromdense(matrix)

args = (csr.data, csr.indices, csr.indptr)
new_csr = CSR(args, shape=(3, 3))  # Wrong shape, does not raise

print(csr, new_csr)
print(CSR.todense(csr), CSR.todense(new_csr))  # Raises error

I think it would be useful to the end user if the error raised in todense would also be raised when creating a new sparse operator.

System info (python version, jaxlib version, accelerator, etc.)

jax: 0.8.1
jaxlib: 0.8.1
numpy: 2.3.3
python: 3.13.3 (main, Apr 8 2025, 13:54:08) [Clang 16.0.0 (clang-1600.0.26.6)]
device info: cpu-1, 1 local devices"
process_count: 1
platform: uname_result(system='Darwin', node=REDACTED, release='24.5.0', version='Darwin Kernel Version 24.5.0: Tue Apr 22 19:54:26 PDT 2025; root:xnu-11417.121.6~2/RELEASE_ARM64_T8112', machine='arm64')

[jakevdp]

Thanks! We're not actively working on jax.experimental.sparse at the moment, so it's unlikely we'll fix issues like this.

For context though, the reason we don't validate these objects on construction is because in the context of pytree flattening within JAX transformations, these containers are often constructed with invalid data. So we push the data validation to the point where the object is used as a CSR matrix. This is a weakness of the pytree-based approach that the current implementation is based on: we could probably do better by redefining sparse objects as hijax types, but that would require a relatively large refactoring.

[johannahaffner]

Thanks! We're not actively working on jax.experimental.sparse at the moment, so it's unlikely we'll fix issues like this.

That is very good to know. Do you have plans to graduate the sparse module out of experimental at some point, or

we could probably do better by redefining sparse objects as hijax types.

does this mean that since there is a better type available (if only in theory) long-term support for sparse data structures is going to be implemented differently in the future?

Also for context, from our side - we do want to add support for large-scale constrained minimisation to Optimistix, backing out to CuDSS for the linear solves. So naturally, the long-term roadmap for sparse data types in JAX is something we're extremely interested in.

Of more practical and immediate concern - I've picked CSR and COO to wrap rather than their batched counterparts, since I don't have a good overview for what BCSR/BCOO will do on top of the sparse representation and find it hard to gauge how this will interact with our filtered transformations. (Vmapping works find out of the box, as do JVPs of our linear solves - but I have not done much sparsity-specific testing.)

So we push the data validation to the point where the object is used as a CSR matrix.

Got it!

[jakevdp]

Do you have plans to graduate the sparse module out of experimental at some point

I don't have any concrete plans to do this at the moment, but we may try it mainly as a way to kick the tires on hijax.

long-term support for sparse data structures is going to be implemented differently in the future?

We don't have any long-term plans for sparse support at the moment. The existing implementations are basically reference implementations, which are backed by very inefficient lowerings to dense gather/scatter ops on the underlying buffers (because the XLA compiler doesn't have any native sparse support). Most feedback indicates the current implementations are too slow to be useful in cases of interest: unless a better sparse-aware backend becomes available, we judged that it was not worth further development on jax.experimental.sparse, and so the submodule is in maintenance mode.

I've picked CSR and COO to wrap rather than their batched counterparts, since I don't have a good overview for what BCSR/BCOO will do on top of the sparse representation

If you ever want to use sparse objects within vmap, or in batched autodiff transformations like jax.jacobian, you'll need to support batched sparse formats. If you don't care about batching, then simple CSR/COO is fine.

[johannahaffner]

Thank you for the detailed response, that is very good to know!

If you ever want to use sparse objects within vmap, or in batched autodiff transformations like jax.jacobian, you'll need to support batched sparse formats. If you don't care about batching, then simple CSR/COO is fine.

Can the operators have different sparsity structures within BCSR/BCOO, provided that the overall operator shape is the same?

---

Our use case right now is limited to interfacing with CuDSS, so all of the interesting/tricky sparse operations actually happen within CuDSS and we just need a nice format to give it the information it needs - the operator in Lineax is a convenience.

I'll keep in mind that any kind of more serious usage of sparsity within JAX will need updated tooling and switching to the batched operators.

[jakevdp]

Can the operators have different sparsity structures within BCSR/BCOO, provided that the overall operator shape is the same?

I don't understand what you're asking.

[johannahaffner]

Indices and indptr are arrays, so I'm wondering if these could have different values as long as they have the same shape across different batch elements. In that case the sparsity structure of the operators thus encoded would be different.

[jakevdp]

The data model is that indices and indptr are batched – so each batch needs to have the same shape, but otherwise the index structure within each batch can differ.

[johannahaffner]

Perfect, thank you for the clarification!