ENH: JAX support in SciPy tracker

[crusaderky]

I'm opening this issue to track what we need in array-api-extra to unblock work in scipy for JAX adoption:

array-api-extra (not blocked by array-api-compat release):

• ENH: Test tools for jax.jit and dask #115
• ENH: at support for bool mask in Dask and JAX #121
• ENH: lazy_apply #86
• ENH: add lazywhere #14
• BUG: vendored xpx.testing crashes on import #127
• ENH: broadcast_shapes #133
• BUG: at should not force overwrite in Dask when copy=None #135

Not on the critical path:

• MAINT: Upgrade to array-api-compat 1.11 #120 (blocked by array-api-compat release)
• BUG: setdiff1d for jax.jit and dask NaN-shaped arrays #116 (only used in scipy.optimize.quadratic_assignment, which is numpy-only)

Shelved indefinitely due to incompatibility with JAX:

• DNM WIP ENH: More lazy functions #99

CC @rgommers

[lucascolley]

Has this issue served its course? Shall we close this and continue tracking at scipy/scipy#22246 ?

[crusaderky]

Yup.