jax

aaac4f93 - Merge pull request #18127 from rwitten:rwitten_make_array_from_single_device_arrays_docs