Jake Vanderplas

One thing I'm unclear on: what is the difference between materialized and non-materialized arrays in terms of the array API? What array API operations can you do on one, but...

OK, thanks for the clarification. In that case, `materialize` is unimplementable in JAX, because materializing an array involves exiting a transformation context, and that can't be done by anything except...

The current error is due to jaxlib 0.4.32 being yanked. If you rebase on the updated main branch, then tests should pass. Thanks!

Regarding reverse-sorting along an axis: does `np.flip(np.sort(x, axis=axis), axis=axis)` do what you want here? Setting aside questions about where `NaN` values end up in reverse-sorted arrays (which is left unspecified...

Thanks! This issue also comes up in a number of other places in this file – would you like to fix those as well as part of this PR?

Regarding the CI failure: in `test_bcoo_dynamic_slice` we need to cast `slice_sizes` to a list. i.e. this: https://github.com/google/jax/blob/c33ce857847750b836b8c899c5d48c12b2842afc/tests/sparse_bcoo_bcsr_test.py#L1062 should be changed to this: ```python slice_sizes = rng.randint(0, shape, len(shape)).tolist() ``` We...

Actually, strike that idea: what we need is to move this line to above the `eval_shape` call inside `bcoo_dynamic_slice`: https://github.com/google/jax/blob/c33ce857847750b836b8c899c5d48c12b2842afc/jax/experimental/sparse/bcoo.py#L2046

No need to compile jaxlib: you should be able to use the most recent jaxlib release along with a local install of `jax` via `pip install -e .[cpu]`

That error generally comes from version skew between jax and jaxlib. Try doing `pip install -U .[cpu]` assuming you're running locally on a CPU, and haven't installed a GPU jaxlib...

Thanks again for this! It was a tricky issue and this fix will make a lot of code more efficient.