openxla / xla

A machine learning compiler for GPUs, CPUs, and ML accelerators
Apache License 2.0
2.39k stars 356 forks source link

Explicitly disallow duplicated devices during array construction #14272

Closed copybara-service[bot] closed 2 days ago

copybara-service[bot] commented 3 days ago

Explicitly disallow duplicated devices during array construction

jax.make_array_from_single_device_arrays should not allow passing more than one array on the same device as that would lead to an invalid array. While some of this case is already detected by later checks (e.g., ArrayImpl._check_and_rearrange), this CL explicitly checks the device list before calling IFRT so that we don't create an invalid IFRT array to begin with.