(x, devices: Sequence[xb.xla_client.Device], replicate: bool = False)
| 62 | |
| 63 | |
| 64 | def device_put(x, devices: Sequence[xb.xla_client.Device], replicate: bool = False): |
| 65 | if replicate: |
| 66 | return list( |
| 67 | it.chain.from_iterable(_device_put_impl(x, device) for device in devices) |
| 68 | ) |
| 69 | else: |
| 70 | return list( |
| 71 | it.chain.from_iterable( |
| 72 | _device_put_impl(val, device) for val, device in safe_zip(x, devices) |
| 73 | ) |
| 74 | ) |
| 75 | |
| 76 | |
| 77 | def get_xla_backend_and_device(device=None) -> Tuple[Backend, Sequence[xc.Device]]: |
no test coverage detected