Return a view or a copy having the given index_map. Data transfers across devices are done on separate streams created internally. To make them asynchronous, transferred data is buffered and reflected to the chunks when necessary. Args: index_map (dict f
(self, index_map: dict[int, Any])
| 320 | return self._to_op_mode(mode) |
| 321 | |
| 322 | def reshard(self, index_map: dict[int, Any]) -> DistributedArray: |
| 323 | """Return a view or a copy having the given index_map. |
| 324 | |
| 325 | Data transfers across devices are done on separate streams created |
| 326 | internally. To make them asynchronous, transferred data is buffered and |
| 327 | reflected to the chunks when necessary. |
| 328 | |
| 329 | Args: |
| 330 | index_map (dict from int to array indices): Indices for the chunks |
| 331 | that devices with designated IDs own. The current index_map of |
| 332 | a distributed array can be obtained from |
| 333 | :attr:`DistributedArray.index_map`. |
| 334 | """ |
| 335 | new_index_map = _index_arith._normalize_index_map( |
| 336 | self.shape, index_map) |
| 337 | if new_index_map == self.index_map: |
| 338 | return self |
| 339 | |
| 340 | old_chunks_map = self._chunks_map |
| 341 | new_chunks_map: dict[int, list[_Chunk]] = {} |
| 342 | |
| 343 | # Set up new_chunks_map compatible with new_index_map |
| 344 | # as placeholders of chunks |
| 345 | for dev, idxs in new_index_map.items(): |
| 346 | new_chunks_map[dev] = [] |
| 347 | |
| 348 | for idx in idxs: |
| 349 | with Device(dev): |
| 350 | dst_shape = _index_arith._shape_after_indexing( |
| 351 | self.shape, idx) |
| 352 | new_chunk = _Chunk.create_placeholder(dst_shape, dev, idx) |
| 353 | new_chunks_map[dev].append(new_chunk) |
| 354 | |
| 355 | self._prepare_comms_and_streams(index_map.keys()) |
| 356 | |
| 357 | # Data transfer from old chunks to new chunks |
| 358 | # TODO: Reorder transfers to minimize latency |
| 359 | |
| 360 | # The current implementation transfers the same data multiple times |
| 361 | # where chunks overlap. This is particularly problematic when matrix |
| 362 | # multiplication is involved, where one block tends to be shared |
| 363 | # between multiple devices |
| 364 | # TODO: Avoid duplicate data transfers |
| 365 | for src_chunk in chain.from_iterable(old_chunks_map.values()): |
| 366 | src_chunk.flush(self._mode) |
| 367 | |
| 368 | if self._mode is not _modes.REPLICA: |
| 369 | src_chunk = src_chunk.copy() |
| 370 | |
| 371 | for dst_chunk in chain.from_iterable(new_chunks_map.values()): |
| 372 | src_chunk.apply_to( |
| 373 | dst_chunk, self._mode, self.shape, |
| 374 | self._comms, self._streams) |
| 375 | |
| 376 | return DistributedArray( |
| 377 | self.shape, self.dtype, new_chunks_map, self._mode, self._comms) |
| 378 | |
| 379 | def get( |