Update jax.experimental.array_api to v2023.12 API

[Micky774]

This issue tracks the changes necessary to adopt the v2023.12 Array API. This was originally mentioned in #18353. Note that there may be some specifications that we already satisfy, however the vast majority of these will need alterations.

API Updates

• Add support for max_version, dl_device, copy kwargs in __dlpack__ to match Array API #20198 | spec
• Add support for device and copy kwargs in from_dlpack to match Array API #20175 | spec
• Update error type in __dlpack__ and __dlpack_device__ to match Array API #20194 | spec
• Add support for copy kwarg in astype to match Array API #20195 | spec

New API

• Add __array_namespace_info__ and corresponding utilities #20294 | spec
  • Add capabilities | spec
  • Add default_device | spec
  • Add default_dtypes | spec
  • Add dtypes | spec
  • Add devices | spec
• Expose existing functions in array API namespace #20753
  • Add copysign | spec
  • Add maximum | spec
  • Add minimum | spec
  • Add moveaxis | spec
  • Add repeat | spec
  • Add searchsorted | spec
  • Add signbit | spec
• Update jnp.clip to Array API 2023 standard and introduces jax.experimental.array_api.clip #20550 | spec
• Add new unstack function to numpy/array_api namespaces #20755 | spec
• Fixed hypot bug on nan/inf pairings, began deprecation of non-real values #20754 | spec
• Add new cumulative_sum function to numpy and array_api namespaces #20756 | spec
• Expose tile function in array_api namespace #20954

Breaking Changes

For specific details on what has changed, look on their specification pages for "Changed in version 2023.12 ..."

• Refactored common upcast for integral-type accumulators  #20842
  • Update prod | spec
  • Update sum | spec

Common Utilities Refactor

• Device placement + copy semantics
  • Initial draft can be found here
  • Get Device from Device | Sharding

[jakevdp]

Thanks for adding this! One overall note: our eventual goal is to remove jax.experimental.array_api and just use the jax.numpy namespace directly. So as much as possible, we should aim for API functions in jax.experimental.array_api to not have any additional logic beyond just calling the jax.numpy counterpart.

[jakevdp]

A note on default_device (cf. a discussion with @yashk2810). JAX doesn't really have a concept of "default device" in the way that the Array API envisions it. By default, arrays are created uncommitted, so the only way to write a consistent default_device function would be for it to return "uncommitted". Currently most (if not all) functions that accept a device parameter have a default value None, so our default_device function should probably look like this:

def default_device():
  return None

and be documented appropriately. That's the only way for, e.g. jax.device_put(x, device=jnp.default_device()) to have the same behavior as jax.device_put(x), which seems like a sensible requirement for the concept of a default!