Broadcast an array to all workers Parameters ---------- src: Union[np.ndarray, Tensor] The array to be broadcasted. dst: Optional[DRef] The output array. If None, an array matching the shape and dtype of `src` will be allocated o
(
self,
src: np.ndarray | Tensor,
dst: DRef | None = None,
in_group: bool = True,
)
| 341 | self._clear_ipc_memory_pool() |
| 342 | |
| 343 | def broadcast( |
| 344 | self, |
| 345 | src: np.ndarray | Tensor, |
| 346 | dst: DRef | None = None, |
| 347 | in_group: bool = True, |
| 348 | ) -> DRef: |
| 349 | """Broadcast an array to all workers |
| 350 | |
| 351 | Parameters |
| 352 | ---------- |
| 353 | src: Union[np.ndarray, Tensor] |
| 354 | The array to be broadcasted. |
| 355 | |
| 356 | dst: Optional[DRef] |
| 357 | The output array. If None, an array matching the shape |
| 358 | and dtype of `src` will be allocated on each worker. |
| 359 | |
| 360 | in_group: bool |
| 361 | Whether the broadcast operation performs globally or in group as default. |
| 362 | |
| 363 | Returns |
| 364 | ------- |
| 365 | output_array: DRef |
| 366 | |
| 367 | The DRef containing the broadcasted data on all workers. |
| 368 | If `dst` was provided, this return value is the same as |
| 369 | `dst`. Otherwise, it is the newly allocated space. |
| 370 | |
| 371 | """ |
| 372 | if not isinstance(src, Tensor): |
| 373 | src = _as_Tensor(src) |
| 374 | |
| 375 | if dst is None: |
| 376 | dst = self.empty(src.shape, src.dtype) |
| 377 | |
| 378 | src_dref = self.copy_to_worker_0(src) |
| 379 | self.broadcast_from_worker0(src_dref, dst, in_group) |
| 380 | |
| 381 | return dst |
| 382 | |
| 383 | def broadcast_from_worker0(self, src: DRef, dst: DRef, in_group: bool = True) -> DRef: |
| 384 | """Broadcast an array from worker-0 to all other workers. |
nothing calls this directly
no test coverage detected