Apply jnp.take with a BCOO array?

[161siegels]

How would i apply jnp.take to a BCOO array? Currently it gives me the following error when i try to call sparse.sparsify(jnp.take)(X,Y):

NotImplementedError: sparse rule for lt is not implemented because it would result in dense output. If this is your intent, use sparse.todense() to convert your arguments to dense matrices.

X is a jnp.array and Y is a BCOO

[jakevdp]

It looks like you're attempting to index a dense array with sparse indices.

Under the assumption that a sparse array is just an array with the zeros not explicitly stored, then the output must be dense (becuse X[0] cannot be guaranteed to be zero). So semantically you can do no better than jnp.take(X, Y.todense()), and that's how I'd recommend solving your problem.