Home
last modified time | relevance | path

Searched refs:new_while (Results 1 – 8 of 8) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/
Dwhile_loop_simplifier_test.cc609 HloInstruction* new_while = FindFirstWhile(m.get()); in TEST_F() local
613 EXPECT_TRUE(ShapeUtil::Equal(new_while->shape(), flat_tuple)); in TEST_F()
615 new_while->while_body()->root_instruction()->shape(), flat_tuple)); in TEST_F()
617 new_while->while_body()->parameter_instruction(0)->shape(), flat_tuple)); in TEST_F()
619 new_while->while_condition()->parameter_instruction(0)->shape(), in TEST_F()
695 HloInstruction* new_while = FindFirstWhile(m.get()); in TEST_F() local
697 EXPECT_TRUE(ShapeUtil::Equal(new_while->shape(), new_while_shape)); in TEST_F()
699 new_while->while_body()->root_instruction()->shape(), new_while_shape)); in TEST_F()
701 new_while->while_body()->parameter_instruction(0)->shape(), in TEST_F()
704 new_while->while_condition()->parameter_instruction(0)->shape(), in TEST_F()
[all …]
Dwhile_util.cc121 HloInstruction* new_while = containing_computation->AddInstruction( in MakeInstructionsLiveIn() local
129 new_while, while_instr->shape().tuple_shapes_size()); in MakeInstructionsLiveIn()
145 result.new_while_instr = new_while; in MakeInstructionsLiveIn()
Dwhile_loop_simplifier.cc1177 auto* new_while = computation->AddInstruction(HloInstruction::CreateWhile( in TryMergeInductionVariables() local
1183 while_op, convert_to_old_form(new_while))); in TryMergeInductionVariables()
1187 return new_while; in TryMergeInductionVariables()
/external/tensorflow/tensorflow/compiler/mlir/xla/transforms/
Dlegalize_tf_communication.cc656 auto new_while = builder.create<WhileOp>(region_while.getLoc(), in RewriteRegionWhileOp() local
660 new_while.cond().takeBody(region_while.cond()); in RewriteRegionWhileOp()
661 new_while.body().takeBody(region_while.body()); in RewriteRegionWhileOp()
666 new_while.getResult()); in RewriteRegionWhileOp()
669 new_while.getLoc(), new_while.getResult(), in RewriteRegionWhileOp()
670 new_while.getResult().getType().cast<TupleType>().size() - 1); in RewriteRegionWhileOp()
681 ops_to_visit.push_back({/*region_idx=*/0, new_token, new_while}); in RewriteRegionWhileOp()
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/
Dtensor_list_ops_decomposition.cc209 auto new_while = in HandleWhileOp() local
213 (*buffer_to_size)[new_while.getResult(std::get<0>(entry))] = { in HandleWhileOp()
214 new_while.getResult(std::get<1>(entry)), std::get<2>(entry)}; in HandleWhileOp()
217 new_while.getResults().take_front(while_op.getNumResults())); in HandleWhileOp()
330 auto new_while = builder.create<TF::WhileRegionOp>( in HandleWhileRegionOp() local
333 new_while.body().takeBody(body_region); in HandleWhileRegionOp()
334 new_while.cond().takeBody(cond_region); in HandleWhileRegionOp()
336 (*buffer_to_size)[new_while.getResult(std::get<0>(entry))] = { in HandleWhileRegionOp()
337 new_while.getResult(std::get<1>(entry)), std::get<2>(entry)}; in HandleWhileRegionOp()
340 new_while.getResults().take_front(while_op.getNumResults())); in HandleWhileRegionOp()
Dstack_ops_decomposition.cc205 auto new_while = in HandleWhileOp() local
220 new_while.getResults().take_front(while_op.getNumResults())); in HandleWhileOp()
Dresource_op_lifting.cc930 auto new_while = builder.create<TF::WhileOp>( in HandleWhileLoop() local
946 AddLoadsStoresOutsideControlFlowOp(new_while, in HandleWhileLoop()
952 new_while.getResult(old_to_new_indices[i])); in HandleWhileLoop()
Dtensor_array_ops_decomposition.cc626 auto new_while = in HandleWhileOp() local
633 while_op.getResult(i).replaceAllUsesWith(new_while.getResult(i)); in HandleWhileOp()