Home
last modified time | relevance | path

Searched refs:scatter_operands (Results 1 – 7 of 7) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/
Dscatter_simplifier.cc160 auto scatter_operands, in ExpandInstruction()
161 MaybeTranspose(scatter->scatter_operands(), operand_permutation)); in ExpandInstruction()
166 if (scatter_operands.size() == 1) { in ExpandInstruction()
167 output_shape = scatter_operands.front()->shape(); in ExpandInstruction()
170 shapes.reserve(scatter_operands.size()); in ExpandInstruction()
171 for (auto* operand : scatter_operands) { in ExpandInstruction()
177 output_shape, scatter_operands, scatter_indices, scatter_updates, in ExpandInstruction()
188 if (scatter->scatter_operands().size() == 1) { in ExpandInstruction()
193 result_items.reserve(scatter->scatter_operands().size()); in ExpandInstruction()
194 for (int i = 0; i < scatter->scatter_operands().size(); ++i) { in ExpandInstruction()
Dscatter_expander.cc415 auto scatter_operands = scatter->scatter_operands(); in ExpandInstruction() local
424 if (scatter_operands.size() == 1) { in ExpandInstruction()
425 return scatter_operands[0]; in ExpandInstruction()
428 HloInstruction::CreateTuple(scatter_operands)); in ExpandInstruction()
469 absl::c_copy(scatter_operands, std::back_inserter(loop_state)); in ExpandInstruction()
483 absl::MakeSpan(scatter_loop_result).first(scatter_operands.size()); in ExpandInstruction()
Dhlo_sharding_util.cc1144 GetScatterSliceSize(scatter.scatter_operands()[0]->shape(), in ScatterOutputShardingFromUpdate()
1163 GetScatterSliceSize(scatter.scatter_operands()[0]->shape(), in ScatterUpdateShardingFromOutput()
1166 scatter.scatter_operands()[0]->shape(), per_output_sharding, in ScatterUpdateShardingFromOutput()
Dhlo_instructions.h2168 absl::Span<HloInstruction* const> scatter_operands() const { in scatter_operands() function
Dalgebraic_simplifier.cc5970 ReplaceInstructionIfCompatible(scatter, scatter->scatter_operands())) { in HandleScatter()
5975 SameShape(scatter, scatter->scatter_operands()[0]) && in HandleScatter()
5979 {scatter->scatter_operands()[0], in HandleScatter()
Dhlo_evaluator.cc2845 for (HloInstruction* operand_inst : scatter->scatter_operands()) { in HandleScatter()
/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dgather_scatter_handler.cc1048 scatter->scatter_operands(), std::back_inserter(operands), in HandleScatter()