Tracking Issue: JAX Type Annotations

[jakevdp]

This is an issue to track progress on JAX Type Annotations.

For information, see the Type Annotation Roadmap, added in #11859

[jakevdp]

jax.Array is lacking a number of attributes that are defined on the array implementation. Here's a list (found by importing the contents of basearray.pyi and running set(dir(jnp.arange(10)) - set(dir(basearray.Array)) and then removing properties and private attributes)

{
 'addressable_data',
 'block_until_ready',
 'broadcast',
 'broadcast_in_dim',
 'copy_to_host_async',
 'delete',
 'device',
 'device_buffer',
 'device_buffers',
 'devices',
 'flat',
 'global_shards',
 'is_deleted',
 'is_fully_addressable',
 'is_fully_replicated',
 'itemsize',
 'on_device_size_in_bytes',
 'split',
 'traceback',
 'unsafe_buffer_pointer'
}

A few of these are obvious omissions (broadcast, broadcast_in_dim, itemsize, split); I've added them in #13967. 'flat' is defined but raises a NotImplementedError on both arrays and tracers; the remaining are array-only attributes and methods that we should probably define on jax.Array and return an informative error in the case of tracers.

[jakevdp]

Closing this, since the core work is complete.