(self, devices: Iterable[int])
| 209 | return chunks_map |
| 210 | |
| 211 | def _prepare_comms_and_streams(self, devices: Iterable[int]) -> None: |
| 212 | # Ensure communicators and streams are prepared for communication |
| 213 | # between `devices` and the devices currently owning chunks |
| 214 | devices = self._chunks_map.keys() | devices |
| 215 | |
| 216 | if not devices.issubset(self._comms.keys()): |
| 217 | self._comms = _data_transfer._create_communicators(devices) |
| 218 | |
| 219 | for dev in devices - self._streams.keys(): |
| 220 | with Device(dev): |
| 221 | self._streams[dev] = Stream() |
| 222 | |
| 223 | def __cupy_override_elementwise_kernel__(self, kernel, *args, **kwargs): |
| 224 | # This method is called from cupy.ufunc and cupy.ElementwiseKernel |
no outgoing calls
no test coverage detected