Searched refs:_build_ring_gather (Results 1 – 2 of 2) sorted by relevance
109 output_tensors = ar._build_ring_gather(input_tensors, device_names, 1,116 output_tensors, pad_len = ar._build_ring_gather(
283 chunks_by_dev, pad_len = _build_ring_gather(297 def _build_ring_gather(input_tensors, devices, num_subchunks, function