Home
last modified time | relevance | path

Searched defs:all_reduce (Results 1 – 18 of 18) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/
Dall_reduce_simplifier.cc39 [this](const HloInstruction* all_reduce) -> int64 { in Run()
80 auto all_reduce = all_reduce_and_group_size.first; in Run() local
Dwhile_loop_all_reduce_code_motion.cc66 MovableAllReduceContext IsAllReduceMovable(HloInstruction* all_reduce, in IsAllReduceMovable()
68 auto all_reduce_is_summation = [](HloInstruction* all_reduce) -> bool { in IsAllReduceMovable()
272 HloComputation* while_body) -> MovableAllReduceContext { in IsAllReduceMovable()
592 HloInstruction* all_reduce = all_reduce_accumulations_pair.first; in Run() local
Dar_crs_combiner.cc96 auto all_reduce = Cast<HloAllReduceInstruction>(hlo); in HasCombinableReplicaGroup() local
543 auto all_reduce = pair.ar; in RewriteGraph() local
Dall_reduce_combiner.cc202 auto channel_id = [](const HloInstruction* all_reduce) { in CreateComputationGroups()
Dall_reduce_combiner_test.cc221 auto all_reduce = b.AddInstruction(HloInstruction::CreateAllReduce( in TEST_F() local
Dbfloat16_propagation_test.cc211 HloInstruction* all_reduce = in TEST_F() local
Dhlo_verifier.cc1519 auto all_reduce = DynCast<HloAllReduceInstruction>(instruction); in VerifyLayoutConstrainedAllReduce() local
/external/tensorflow/tensorflow/python/ops/
Dcollective_ops.py23 def all_reduce(t, function
/external/tensorflow/tensorflow/python/kernel_tests/
Dcollective_ops_test.py47 all_reduce = _collective_ops.all_reduce variable in CollectiveOpsV1
56 def all_reduce(t, group_size, group_key, instance_key, *args, **kwargs): member in CollectiveOpsV2
/external/tensorflow/tensorflow/python/distribute/
Dcross_device_utils.py335 def all_reduce(self, member in CollectiveReplicaLauncher
Dvars_test.py957 def all_reduce(): function
Dtpu_strategy_test.py639 def all_reduce(x): function
Ddistribute_lib.py3053 def all_reduce(self, reduce_op, value, options=None): member in ReplicaContextBase
/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dgather_scatter_handler.cc605 auto all_reduce = operand.state().partitioner->AllReduceAlongShardingDims( in HandleScatter() local
Dspmd_partitioner.cc1561 auto all_reduce = per_group_partitioner_state.collective_ops_creator in HandleConcatenate() local
/external/tensorflow/tensorflow/compiler/mlir/xla/
Dhlo_function_importer.cc555 auto all_reduce = Cast<HloAllReduceInstruction>(instruction); in ImportInstructionImpl() local
/external/tensorflow/tensorflow/compiler/mlir/xla/transforms/
Dmhlo_to_lhlo_with_xla.cc1079 auto* all_reduce = xla::Cast<xla::HloAllReduceInstruction>(instr); in EmitAllReduceOp() local
Dlegalize_tf.cc5274 auto all_reduce = rewriter.create<AllReduceOp>( in matchAndRewrite() local