Searched refs:reduce_tensor (Results 1 – 4 of 4) sorted by relevance
| /aosp_15_r20/external/pytorch/torch/multiprocessing/ |
| H A D | reductions.py | 223 def reduce_tensor(tensor): function
421 rebuild_buffer_func, rebuild_buffer_args = reduce_tensor(nt.values())
422 rebuild_sizes_func, rebuild_sizes_args = reduce_tensor(nt._nested_tensor_size())
423 rebuild_strides_func, rebuild_strides_args = reduce_tensor(
426 rebuild_offsets_func, rebuild_offsets_args = reduce_tensor(
480 rebuild_indices_func, rebuild_indices_args = reduce_tensor(sparse._indices())
481 rebuild_values_func, rebuild_values_args = reduce_tensor(sparse._values())
505 ) = reduce_tensor(compressed_indices)
506 rebuild_plain_indices_func, rebuild_plain_indices_args = reduce_tensor(
509 rebuild_values_func, rebuild_values_args = reduce_tensor(sparse.values())
[all …]
|
| /aosp_15_r20/external/pytorch/torch/csrc/distributed/c10d/ |
| H A D | default_comm_hooks.cpp | 34 auto reduce_tensor = result.toTensorVector()[0]; in runHook() local
36 reduce_tensor.scalar_type() == at::ScalarType::Half, in runHook()
38 reduce_tensor.scalar_type()); in runHook()
39 decompressed_tensor.copy_(reduce_tensor); in runHook()
|
| /aosp_15_r20/external/pytorch/test/distributed/rpc/ |
| H A D | test_share_memory.py | 46 self._dispatch_table[t] = TorchMpReductions.reduce_tensor
47 self._dispatch_table[torch.Tensor] = TorchMpReductions.reduce_tensor
50 ] = TorchMpReductions.reduce_tensor
|
| /aosp_15_r20/external/pytorch/test/distributed/ |
| H A D | test_c10d_nccl.py | 781 reduce_tensor = torch.rand(10, 10, device=device)
783 pg.allreduce(reduce_tensor).wait()
805 reduce_tensor = torch.rand(10, 10, device=device)
808 pg.allreduce(reduce_tensor).wait()
|