Home
last modified time | relevance | path

Searched refs:while_init (Results 1 – 9 of 9) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/
Dwhile_loop_simplifier.cc56 HloInstruction* while_init = while_op->mutable_operand(0); in TryRemoveDeadWhileParams() local
61 if (!while_init->shape().IsTuple()) { in TryRemoveDeadWhileParams()
89 const int64 tuple_size = ShapeUtil::TupleElementCount(while_init->shape()); in TryRemoveDeadWhileParams()
174 while_init->shape().tuple_shapes(old_idx)); in TryRemoveDeadWhileParams()
257 while_init->shape().tuple_shapes(old_idx), while_init, old_idx))); in TryRemoveDeadWhileParams()
298 while_init->shape().tuple_shapes(old_idx), while_init, old_idx))); in TryRemoveDeadWhileParams()
313 auto* while_init = while_op->mutable_operand(0); in TryRemoveConstantParams() local
317 if (while_init->opcode() != HloOpcode::kTuple || in TryRemoveConstantParams()
325 ShapeUtil::Compatible(while_init->shape(), while_body_root->shape())); in TryRemoveConstantParams()
328 const auto& while_shape = while_init->shape(); in TryRemoveConstantParams()
[all …]
Dwhile_loop_analysis.cc156 auto* while_init = while_op->operand(0); in GetLoopInductionVarTupleIdx() local
157 if (while_init->opcode() != HloOpcode::kTuple) { in GetLoopInductionVarTupleIdx()
158 VLOG(2) << "While init expected to be a tuple: " << while_init->ToString(); in GetLoopInductionVarTupleIdx()
347 auto* while_init = while_op->mutable_operand(0); in ComputeWhileLoopTripCount() local
348 auto* indvar_init = while_init->mutable_operand(*indvar_tuple_idx); in ComputeWhileLoopTripCount()
Dhlo_memory_scheduler_test.cc275 HloInstruction* while_init = in TEST_F() local
280 r1f32, cond_computation, body_computation, while_init)); in TEST_F()
Dcopy_insertion.cc277 HloInstruction* while_init = xla_while->mutable_operand(0); in AddCopiesForWhile() local
280 xla_while->parent()->DeepCopyInstruction(while_init, &indices_to_copy)); in AddCopiesForWhile()
281 TF_RETURN_IF_ERROR(while_init->ReplaceUseWith(xla_while, while_init_copy)); in AddCopiesForWhile()
Dar_crs_combiner.cc123 auto init_tuples = GetAllTuples(while_instr->while_init()); in GetAllTuples()
Dheap_simulator_test.cc167 HloInstruction* while_init = in TEST_F() local
173 r1f32, cond_computation, body_computation, while_init)); in TEST_F()
197 std::vector<HloInstruction*> entry_comp_vec = {while_init, while_loop, bcast, in TEST_F()
Dcopy_insertion_test.cc1753 auto while_init = builder.AddInstruction(HloInstruction::CreateTuple( in TEST_F() local
1757 loop_state_shape, condition, body, while_init)); in TEST_F()
Dhlo_instruction.h1069 HloInstruction* while_init() const;
Dhlo_instruction.cc2052 HloInstruction* HloInstruction::while_init() const { in while_init() function in xla::HloInstruction