RFC: In-memory sparse array interchange

[hameerabbasi]

Motivation

The sparse (also called PyData/Sparse) development team have been working on integration efforts with the ecosystem, most notably with SciPy, scikit-learn and others, with CuPy, PyTorch, JAX and TensorFlow also on the radar. One of the challenges we were facing was the lack of (possibly zero-copy) interchange between the different sparse array implementations. We believe this may be a pain point for many sparse array implementations moving forward.

This mirrors an issue seen for dense arrays previously, where the DLPack protocol was the one of the first things to be standardised. We're hoping to achieve community consensus for a similar problem.

Luckily, all sparse array formats (with the possible exception of DOK) are usually collections of dense arrays underneath. In addition, this problem has been solved for on-disk arrays before by the binsparse specification. @willow-ahrens is a co-author of that spec, and is also a collaborator for the sparse work.

Proposal

We propose introducing two new methods to the array-API compliant sparse array objects (such as those in sparse), which are described below.

__binsparse_descriptor__

The returned item is a dict equivalent to a parsed JSON binsparse descriptor of an array.

__binsparse__

The second item is a dict[str, Array] of __dlpack__ compatible arrays, which are the constituent arrays of the sparse array. The key represents the equivalent key in the descriptor.

Introduction of from_binsparse function.

If a library supports sparse arrays, its from_binsparse method should support accepting (when possible, zero-copy) versions of objects that follow this __binsparse__ protocol, and have an equivalent sparse format within the library.

Psuedocode implementation

Here's a psuedocode example using two libraries, xp1 and xp2, both supporting sparse arrays:

# In library code:
xp2_sparray = xp2.from_binsparse(xp1_sparray, ...)

# This psuedocode impl is common between `xp1` and `xp2`
def from_binsparse(x: object, /, *, device: device | None = None, copy: bool | None = None) -> array:
    binsparse_descr = getattr(x, "__binsparse_descriptor__", None)
    binsparse_impl = getattr(x, "__binsparse__", None)
    if binsparse_impl is None or binsparse_descr is None:
        raise TypeError(...)
    
    binsparse_descriptor = binsparse_descr()
    # Will raise an error if the format/descriptor is unsupported.
    sparse_type = _type_from_binsparse_descriptor(binsparse_descriptor)
    constituent_arrays = binsparse_impl()
    my_constituent_arrays = {
        k: from_dlpack(arr, device=device, copy=copy) for k, arr in constituent_arrays.items()
    }
    return sparse_type.from_strided_arrays(my_constituent_arrays, shape=...)

Parallel implementation in sparse: pydata/sparse#764
Parallel implementation in SciPy: scipy/scipy#22553

Alternative solutions

There are formats for on-disk sparse-array interchange [1] [2]; but none for in-memory interchange. binsparse is the one that comes closest to offering in-memory interchange.

Pinging possibly interested parties:

• @mtsokol @ivirshup (for scipy.sparse)
• @willow-ahrens @mtsokol (from binsparse and finch-tensor/sparse)
• @leofang (for cupyx.sparse)
• @pearu (for torch.sparse)
• @jakevdp (for JAX/TensorFlow)

Updated on 2024.10.09 as agreed in #840 (comment).

[pearu]

A couple of quick notes.

First, dlpack and binsparse are self-contained specifications where dlpack provides a protocol for sharing strided arrays between different array libraries and binsparse provides a unified description of various sparse formats. This proposal tries to glue these together by implicitly extending DLPack protocol but that contradicts with the semantics of DLPack that "describes the memory layout of dense, strided, n-dimensional arrays". In addition, the Python Array API standard specifies that DLPack should not support sparse arrays.

My suggestion is to cook up a sparse array interchange protocol that may use dlpack and binsparse (as these are obviously relevant pieces of this problem) but in a cleaner way.

For instance, binsparse specifies that the binsparse-compatible object provides a 2-tuple (<sparse array descriptor>, <list of sparse array data as strided arrays>). So, the appropriate place to deploy the dlpack protocol is to require that the arrays in the <list of sparse array data as strided arrays> must support the dlpack protocol.

For instance 2, introduce "Python specification of binsparse" (similar to one in dlpack) that consists of

1. a from_binsparse function that accepts objects that implement the binsparse protocol per 2. below
2. a __binsparse__ method on the objects representing sparse arrays that will return a 2-tuple (<sparse array descriptor>, <list of sparse array data as strided arrays>) which is used by the from_binsparse function of the consumer library.

Also, the specification should include the C-side interface as well.

What do you think?

[hameerabbasi]

While I'm on board with the proposed __binsparse__ protocol, I'm of the opinion that sparse arrays should have as little user-facing API differences as possible as compared to strided arrays, i.e. using sparse arrays should be as frictionless from the Array API side as using strided arrays. Allocating a different function based on sparse and dense arrays (from_dlpack vs from_binsparse) seems counter to that.

I'm not opposed to a C-side specification, I'd welcome that so that the benefits can extend beyond the Python ecosystem. We do have to be mindful that the binsparse specification requires key-value maps, and we would need to maintain their relative order. Perhaps an iterable of 2-tuples makes more sense in this case.

[pearu]

While I'm on board with the proposed __binsparse__ protocol, I'm of the opinion that sparse arrays should have as little user-facing API differences as possible as compared to strided arrays, i.e. using sparse arrays should be as frictionless from the Array API side as using strided arrays.

Absolutely. I presume that users of sparse arrays are aware of advantages and disadvantages of using a particular sparse format over other sparse formats or strided arrays. Using a particular sparse format is a choice of optimization method that must be supported by user-facing API.

Allocating a different function based on sparse and dense arrays (from_dlpack vs from_binsparse) seems counter to that.

from_dlpack is specifically defined for strided arrays, both by the Python Array API and dlpack. So, I find abusing the from_dlpack function for sparse arrays more confusing than providing a new API function which would follow the same design pattern that dlpack introduced: the pair (from_foo, __foo__) provides a minimal API for sharing data between provider and consumer libraries that support a protocol foo.

Btw, if one considers strided and sparse arrays semantically equivalent (as in PyTorch, for instance), a more intuitive approach would be to use asarray instead of from_dlpack as a common API function for constructing strided or sparse arrays from provider library objects. Although, Python Array API does not support that and existing implementations (PyTorch, JAX, ...) do not mix the functionalities of asarray and from_dlpack. I don't suggest to do that here either.

[hameerabbasi]

Btw, if one considers strided and sparse arrays semantically equivalent (as in PyTorch, for instance), a more intuitive approach would be to use asarray instead of from_dlpack as a common API function for constructing strided or sparse arrays from provider library objects.

While I agree with a common API; the disadvantage I find in this approach is that since from_dlpack requires the producer to pass ownership to the consumer; asarray requires one to maintain the ownership of the consumed array. This, in turn, means that there's no way to have a zero-copy transfer using asarray using the from_dlpack function; i.e. one would have to do something like def asarray(x, ...): return from_dlpack(copy(x), ...) somewhere to maintain ownership semantics, incurring a copy. See #840 (comment)

My intention was to have an API that essentially could support zero-copy interchange across libraries, regardless of whether the consumed array was strided, sparse or something else.

[pearu]

While I agree with a common API; the disadvantage I find in this approach is that since from_dlpack requires the producer to pass ownership to the consumer;

This is incorrect. When using from_dlpack(x), the producer keeps owning the memory of x. See https://dmlc.github.io/dlpack/latest/python_spec.html#semantics.

[hameerabbasi]

While I agree with a common API; the disadvantage I find in this approach is that since from_dlpack requires the producer to pass ownership to the consumer;

This is incorrect. When using from_dlpack(x), the producer keeps owning the memory of x. See https://dmlc.github.io/dlpack/latest/python_spec.html#semantics.

Ah, in that case; yes, asarray can use from_dlpack and a possible from_binsparse, and still be zero-copy. Thanks for the correction; asarray seems like the right format-agnostic API for this.

[hameerabbasi]

I've updated the issue with @pearu's feedback.