Searched refs:while_op (Results 1 - 6 of 6) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/
H A Dwhile_loop_simplifier.cc126 const HloInstruction* while_op) {
127 CHECK_EQ(while_op->opcode(), HloOpcode::kWhile);
129 << while_op->ToShortString();
137 auto* while_cond = while_op->while_condition();
155 auto* while_body = while_op->while_body();
183 auto* while_init = while_op->operand(0);
195 static optional<int64> GetLoopTripCount(HloInstruction* while_op) { argument
196 CHECK_EQ(while_op->opcode(), HloOpcode::kWhile);
197 VLOG(2) << "Getting trip count for loop " << while_op->ToString();
203 // where comp is while_op
125 GetLoopInductionVarTupleIdx( const HloInstruction* while_op) argument
286 TryRemoveDeadWhileParams(HloInstruction* while_op) argument
560 TryRemoveWhileLoop(HloInstruction* while_op) argument
[all...]
H A Dhlo_scheduling_test.cc74 HloInstruction* while_op = builder.AddInstruction(HloInstruction::CreateWhile( local
87 module_sequence[entry_computation] = {iter, data, tuple, while_op};
H A Dwhile_loop_simplifier_test.cc128 auto* while_op = computation->root_instruction(); local
129 ASSERT_EQ(while_op->opcode(), HloOpcode::kWhile);
130 auto* true_op = while_op->while_body()->AddInstruction(
133 while_op->while_body()->root_instruction()));
144 auto* while_op = computation->root_instruction(); local
145 ASSERT_EQ(while_op->opcode(), HloOpcode::kWhile);
146 auto* while_body = while_op->while_body();
157 auto* while_op = computation->root_instruction(); local
158 ASSERT_EQ(while_op->opcode(), HloOpcode::kWhile);
159 auto* while_body = while_op
173 auto* while_op = computation->root_instruction(); local
391 auto* while_op = builder.AddInstruction(HloInstruction::CreateWhile( local
[all...]
H A Dheap_simulator_test.cc442 HloInstruction* while_op = builder.AddInstruction(HloInstruction::CreateWhile( local
447 {param, while_op, body_param, cond_param, cond_iter, cond_data, cond_lt});
449 // The entry computation param and while_op are allocated first.
453 {kAlloc, tracker.BufferAt(while_op, {})},
454 {kAlloc, tracker.BufferAt(while_op, {0})},
455 {kAlloc, tracker.BufferAt(while_op, {1})},
481 {kFree, tracker.BufferAt(while_op, {})},
482 {kFree, tracker.BufferAt(while_op, {0})},
483 {kFree, tracker.BufferAt(while_op, {1})},
H A Dbuffer_assignment_test.cc701 auto while_op = builder.AddInstruction(HloInstruction::CreateWhile( local
705 const std::vector<const HloInstruction*> level0 = GetInstructions(while_op);
735 // Check that buffer for each subshape of 'while_op' shares allocation with
738 while_op->shape(),
739 [this, &buffers, while_op, body_root](const Shape& /*subshape*/,
741 auto while_op_allocation = GetAllocation(*buffers, while_op, index);
/external/tensorflow/tensorflow/compiler/tf2xla/
H A Dfunctionalize_control_flow_test.cc251 auto while_op = local
254 auto sink = ops::Identity(scope.WithOpName("sink"), while_op[0]);
354 auto while_op = local
496 auto while_op = local
499 auto sink_x = ops::Identity(scope.WithOpName("sink_x"), while_op[0]);
500 auto sink_y = ops::Identity(scope.WithOpName("sink_y"), while_op[1]);
736 auto while_op = ops::XlaWhile(scope.WithOpName("outer/LoopCond"), local
739 auto sink = ops::Identity(scope.WithOpName("sink"), while_op[0]);
792 auto while_op = local
802 while_op[
[all...]

Completed in 192 milliseconds