Home
last modified time | relevance | path

Searched refs:loop_state_shape (Results 1 – 3 of 3) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/
Dwhile_util.cc154 MakeCountedLoopConditionComputation(const Shape& loop_state_shape, in MakeCountedLoopConditionComputation() argument
160 {&loop_state_shape}, scalar_pred, "while_cond")); in MakeCountedLoopConditionComputation()
177 const Shape& loop_state_shape, in MakeCountedLoopBodyComputation() argument
182 {&loop_state_shape}, loop_state_shape, "while_body")); in MakeCountedLoopBodyComputation()
192 for (int64 i = 1, e = loop_state_shape.tuple_shapes_size(); i < e; i++) { in MakeCountedLoopBodyComputation()
249 Shape loop_state_shape = MakeLoopStateShapeWithLayout(init_values); in MakeCountedLoop() local
252 MakeCountedLoopConditionComputation(loop_state_shape, trip_count)); in MakeCountedLoop()
255 MakeCountedLoopBodyComputation(loop_state_shape, loop_body_generator)); in MakeCountedLoop()
261 loop_state_shape, module->AddEmbeddedComputation(std::move(cond)), in MakeCountedLoop()
Dcopy_insertion_test.cc414 const Shape& loop_state_shape) { in BuildConditionComputation() argument
419 HloInstruction::CreateParameter(0, loop_state_shape, "loop_state")); in BuildConditionComputation()
476 const Shape& loop_state_shape = ShapeUtil::MakeTupleShape( in BuildDependentBodyComputation2() local
480 HloInstruction::CreateParameter(0, loop_state_shape, "loop_state")); in BuildDependentBodyComputation2()
549 const Shape& loop_state_shape = in BuildIndependentBodyComputation() local
553 HloInstruction::CreateParameter(0, loop_state_shape, "loop_state")); in BuildIndependentBodyComputation()
760 const Shape& loop_state_shape, HloInstruction* data_init, in BuildWhileInstructionWithCustomInit() argument
763 ShapeUtil::Equal(loop_state_shape, nested_loop_state_shape_); in BuildWhileInstructionWithCustomInit()
767 BuildConditionComputation(loop_state_shape)); in BuildWhileInstructionWithCustomInit()
773 loop_state_shape, condition, body, loop_state_init)); in BuildWhileInstructionWithCustomInit()
[all …]
/external/tensorflow/tensorflow/compiler/xla/tests/
Dwhile_test.cc1272 Shape loop_state_shape = ShapeUtil::MakeTupleShape( in BM_WhileLoop() local
1281 auto prev = Parameter(&builder, 0, loop_state_shape, "prev"); in BM_WhileLoop()
1291 auto prev = Parameter(&builder, 0, loop_state_shape, "prev"); in BM_WhileLoop()