Searched defs:src_rank (Results 1 – 7 of 7) sorted by relevance
| /aosp_15_r20/external/pytorch/torch/distributed/_shard/ |
| H A D | api.py | 17 tensor: torch.Tensor, sharding_spec: ShardingSpec, src_rank=0, process_group=None argument
81 src_rank=0, argument
226 def shard_module(module: nn.Module, plan: ShardingPlan, src_rank=0, process_group=None): argument
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/common_runtime/ |
| H A D | hierarchical_tree_broadcaster.cc | 45 string BroadcastBufKey(const string& exec_key, int subdiv, int src_rank, in BroadcastBufKey()
408 int src_rank, in DispatchSend()
433 void HierarchicalTreeBroadcaster::DispatchRecv(int subdiv, int src_rank, in DispatchRecv()
|
| H A D | permuter.cc | 80 void Permuter::DispatchSend(int src_rank, int target_rank, const Tensor* tensor, in DispatchSend()
96 void Permuter::DispatchRecv(int src_rank, int target_rank, Tensor* tensor, in DispatchRecv()
|
| H A D | all_to_all.cc | 123 void AllToAll::DispatchSend(int src_rank, int target_rank, const Tensor* tensor, in DispatchSend()
135 void AllToAll::DispatchRecv(int src_rank, int target_rank, Tensor* tensor, in DispatchRecv()
|
| H A D | eval_const_tensor.cc | 87 int src_rank = shape_c->Rank(shape_c->input(0)); in TryToInferTensorOutputFromShapeNode() local
|
| /aosp_15_r20/external/pytorch/torch/csrc/distributed/c10d/ |
| H A D | CUDASymmetricMemory.cu | 144 int src_rank = (rank + world_size - step) % world_size; in all_gather_fds() local
154 int src_rank, in broadcast_fds()
472 int src_rank, in wait_signal_kernel()
482 void CUDASymmetricMemory::wait_signal(int src_rank, int channel) { in wait_signal()
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/nccl/ |
| H A D | nccl_manager_test.cc | 211 TensorShape shape, int src_node, int src_rank, in MakeBroadcastTestCase()
624 const int src_rank = static_cast<int>(random::New64() % num_ranks); in TYPED_TEST() local
|