Home
last modified time | relevance | path

Searched refs:flatOutputTensor (Results 1 – 2 of 2) sorted by relevance

/aosp_15_r20/external/pytorch/torch/csrc/distributed/c10d/
H A DProcessGroupMPI.cpp503 auto flatOutputTensor = newLikeFlat(outputDataVec); in allgather() local
511 flatOutputTensor.data_ptr(), in allgather()
517 outputDataVec[i].copy_(flatOutputTensor[static_cast<int64_t>(i)]); in allgather()
565 at::Tensor flatOutputTensor; in gather() local
569 flatOutputTensor = newLikeFlat(dstdata); in gather()
570 recvbuf = flatOutputTensor.data_ptr(); in gather()
590 flatOutputTensor[static_cast<int64_t>(i)]); in gather()
H A DProcessGroupGloo.cpp1835 at::Tensor flatOutputTensor = newLikeFlat(outputs[0]); in allgather() local
1836 GENERATE_ALL_TYPES(scalarType, setOutput, opts, flatOutputTensor); in allgather()
1842 outputgroup[j].copy_(flatOutputTensor[static_cast<int64_t>(j)]); in allgather()
2114 at::Tensor flatOutputTensor = in allgather_coalesced() local
2116 GENERATE_ALL_TYPES(scalarType, setOutput, opts, flatOutputTensor); in allgather_coalesced()
2123 flatOutputTensor.narrow(0, current_element, output_tensor.numel()) in allgather_coalesced()
2248 at::Tensor flatOutputTensor; in gather() local
2250 flatOutputTensor = newLikeFlat(outputs[0]); in gather()
2251 GENERATE_ALL_TYPES(scalarType, setOutput, opts, flatOutputTensor); in gather()
2261 outputs[0][i].copy_(flatOutputTensor[static_cast<int64_t>(i)]); in gather()