flax

b71bbd66 - Add `make_array_from_single_device_arrays` to prepare to rename of the concrete `Array` to `ArrayImpl`.