Tracking issue: NumPy 2.0 Compatibility

[jakevdp]

NumPy 2.0 is coming, and there are a number of things we need to do to ensure a smooth transition for users. This issue tracks these TODOs. Relevant NumPy issue is here: numpy/numpy#24300.

This has some overlap with #18353, as NumPy is aiming for array API compatibility in v2.0.

• add an upstream nightly build (kokoro only, because of jaxlib build requirement)
• fix np.ComplexWarning references (Remove reference to np.ComplexWarning #19245)
• new sign convention for complex entries (jnp.sign: use x/abs(x) for complex arguments #19390)
• fix jax.scipy.special.logsumexp for new sign convention (logsumexp: use NumPy 2.0 convention for complex sign #19389)
• New shape for inverse_indices in jnp.unique (jnp.unique: make return_inverse shape match NumPy 2.0 #19320)
• add new jax.numpy functions:
  • concat [array api] add jax.numpy.concat #19323
  • isdtype Add jnp.isdtype function, following np.isdtype in NumPy 2.0 #19400
  • permute_dims [array api] add jax.numpy.permute_dims function #19244
  • bitwise_invert array api: add jnp.bitwise_* aliases #19278
  • bitwise_left_shift array api: add jnp.bitwise_* aliases #19278
  • bitwise_right_shift array api: add jnp.bitwise_* aliases #19278
  • pow [array API] implement jnp.pow; alias for jnp.power #19293
  • vecdot Add jnp.vecdot #19274

Fix some relevant bugs in NumPy & SciPy:

• BUG: Fix return shape of inverse_indices in unique_inverse numpy/numpy#25553
• BUG: support axes argument in np.linalg.tensordot numpy/numpy#25554
• BUG: np.linalg.vector_norm: return correct shape for keepdims numpy/numpy#25560
• MAINT: logsumexp: properly handle complex sign scipy/scipy#19870

Once there is a NumPy 2.0 release candidate, we should do an ml_dtypes and jaxlib release built against that candidate; these should be compatible with numpy versions back to our minimum supported version.

[jakevdp]

We're in good shape here; the only failures as of today are fft bugs upstream (numpy/numpy#25661 and numpy/numpy#25679). These are being fixed in numpy/numpy#25668

Beyond that, our main TODO is to update our builds to use the numpy 2.0 ABI once the release candidates are out.

[hawkinsp]

I think we can declare this fixed. We've made a release with NumPy 2.0, which was the last main TODO.