| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
| H A D | all_reduce_blueconnect_test.cc | 77 auto all_gather = AllOf(op::Shape("f32[16]"), op::AllGather(all_reduce), in TEST_F() local
166 auto all_gather = AllOf(op::Shape("(f32[16], f32[32])"), in TEST_F() local
|
| H A D | all_reduce_blueconnect.cc | 239 HloInstruction* all_gather = in TryDecomposeAllReduce() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/utils/ |
| H A D | collective_lowering.cc | 442 mlir::LogicalResult LowerAllGatherOp(mlir::TF::DTensorAllGatherOp all_gather) { in LowerAllGatherOp()
778 module.walk([&](mlir::TF::DTensorAllGatherOp all_gather) { in runOnOperation()
782 for (mlir::TF::DTensorAllGatherOp all_gather : all_gathers) in runOnOperation() local
|
| /aosp_15_r20/external/pytorch/torch/distributed/ |
| H A D | collective_utils.py | 105 def all_gather( function
|
| H A D | distributed_c10d.py | 3275 def all_gather(tensor_list, tensor, group=None, async_op=False): function
|
| /aosp_15_r20/external/pytorch/torch/cuda/ |
| H A D | nccl.py | 131 def all_gather( function
|
| /aosp_15_r20/external/tensorflow/tensorflow/python/kernel_tests/ |
| H A D | collective_ops_test.py | 53 def all_gather(t, group_size, group_key, instance_key, *args, **kwargs): member in CollectiveOpsV1
73 def all_gather(t, group_size, group_key, instance_key, *args, **kwargs): member in CollectiveOpsV2
|
| /aosp_15_r20/external/pytorch/test/distributed/tensor/parallel/ |
| H A D | test_tp_examples.py | 47 reduce_scatter, all_gather, all_reduce = ( variable
|
| /aosp_15_r20/external/tensorflow/tensorflow/python/ops/ |
| H A D | collective_ops.py | 159 def all_gather(t, function
|
| /aosp_15_r20/external/tensorflow/tensorflow/python/distribute/ |
| H A D | cross_device_utils.py | 434 def all_gather( member in CollectiveReplicaLauncher
|
| H A D | tpu_strategy.py | 1726 def all_gather(self, value, axis, experimental_hints=None): member in _TPUReplicaContext
|
| H A D | distribute_lib.py | 3302 def all_gather(self, value, axis, options=None): member in ReplicaContext
|
| /aosp_15_r20/external/pytorch/torch/distributed/checkpoint/ |
| H A D | utils.py | 234 def all_gather( member in _DistWrapper
|
| /aosp_15_r20/external/pytorch/test/inductor/ |
| H A D | test_distributed_patterns.py | 20 def all_gather(t): function
|
| /aosp_15_r20/external/pytorch/torch/distributed/nn/ |
| H A D | functional.py | 107 def all_gather(tensor, group=group.WORLD): function
|
| /aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/ |
| H A D | collectives.cc | 93 mlir::TF::DTensorAllGatherOp all_gather = in EmitAllGather() local
|
| /aosp_15_r20/external/pytorch/test/distributed/_composable/fsdp/ |
| H A D | test_fully_shard_comm.py | 151 def all_gather(fsdp_param_group: FSDPParamGroup, group: dist.ProcessGroup): function
|
| /aosp_15_r20/external/pytorch/torch/csrc/cuda/ |
| H A D | nccl.cpp | 760 void all_gather( in all_gather() function
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/xla/ |
| H A D | hlo_function_importer.cc | 973 auto all_gather = Cast<HloAllGatherInstruction>(instruction); in ImportInstructionImpl() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/xla/transforms/ |
| H A D | mhlo_to_lhlo_with_xla.cc | 1161 auto* all_gather = xla::Cast<xla::HloAllGatherInstruction>(instr); in EmitAllGatherOp() local
|