Batch CSR-CSC matrices?

[ricardoV94]

Here is a gist with a naive implementation of a batch CSR/CSC matrices, which is built by concatenating the data buffers of individuals matrices and keeping an extra piece of info: cum_nnz: https://gist.github.com/ricardoV94/4e57e311db860bf5264267a86553e17c

Recreating the individual matrices is then just a question of slicing the data buffers correctly. Operations on these are just a question of working with the sliced data buffers then. This should be pretty efficient in numba?

Of course half of the challenge of working with sparse types is to create them efficiently (and auto-diffing efficiently), but having a representation to begin with would be great so we can start vectorizing these graphs.

Related to #839

It should be possible to allow arbitrary dimensions, we just need to be clever about how we build and access the cum_nnz based on the batch dims.

We could also have a concept of batch strides to slice across batch dimension without having to copy the data buffers. So something like batch_csr[2:] or batch_csr[::-1] just plays with the pointer/strides in conjunction with cum_nnz.