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.cc152 MakeCountedLoopConditionComputation(const Shape& loop_state_shape, in MakeCountedLoopConditionComputation() argument
158 {&loop_state_shape}, scalar_pred, "while_cond")); in MakeCountedLoopConditionComputation()
175 const Shape& loop_state_shape, in MakeCountedLoopBodyComputation() argument
180 {&loop_state_shape}, loop_state_shape, "while_body")); in MakeCountedLoopBodyComputation()
190 for (int64 i = 1, e = loop_state_shape.tuple_shapes_size(); i < e; i++) { in MakeCountedLoopBodyComputation()
234 Shape loop_state_shape = MakeLoopStateShape(init_values); in MakeCountedLoop() local
237 MakeCountedLoopConditionComputation(loop_state_shape, trip_count)); in MakeCountedLoop()
240 MakeCountedLoopBodyComputation(loop_state_shape, loop_body_generator)); in MakeCountedLoop()
246 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()
473 const Shape& loop_state_shape = ShapeUtil::MakeTupleShape( in BuildDependentBodyComputation2() local
477 HloInstruction::CreateParameter(0, loop_state_shape, "loop_state")); in BuildDependentBodyComputation2()
543 const Shape& loop_state_shape = in BuildIndependentBodyComputation() local
547 HloInstruction::CreateParameter(0, loop_state_shape, "loop_state")); in BuildIndependentBodyComputation()
754 const Shape& loop_state_shape, HloInstruction* data_init, in BuildWhileInstructionWithCustomInit() argument
757 ShapeUtil::Equal(loop_state_shape, nested_loop_state_shape_); in BuildWhileInstructionWithCustomInit()
761 BuildConditionComputation(loop_state_shape)); in BuildWhileInstructionWithCustomInit()
767 loop_state_shape, condition, body, loop_state_init)); in BuildWhileInstructionWithCustomInit()
[all …]
/external/tensorflow/tensorflow/compiler/xla/tests/
Dwhile_test.cc1273 Shape loop_state_shape = ShapeUtil::MakeTupleShape( in BM_WhileLoop() local
1282 auto prev = Parameter(&builder, 0, loop_state_shape, "prev"); in BM_WhileLoop()
1292 auto prev = Parameter(&builder, 0, loop_state_shape, "prev"); in BM_WhileLoop()