jax

a99ea733 - Use `jax.make_array_from_process_local_data` API in distributed data loading doc