Home
last modified time | relevance | path

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

/external/tensorflow/tensorflow/compiler/xla/service/
Dwhile_util.cc178 TF_ASSIGN_OR_RETURN(std::unique_ptr<HloComputation> body_computation, in MakeCountedLoopBodyComputation()
181 HloInstruction* one = body_computation->AddInstruction( in MakeCountedLoopBodyComputation()
183 HloInstruction* param = body_computation->parameter_instruction(0); in MakeCountedLoopBodyComputation()
199 body_computation->AddInstruction(HloInstruction::CreateTuple(next_state)); in MakeCountedLoopBodyComputation()
200 body_computation->set_root_instruction(next_state_tuple); in MakeCountedLoopBodyComputation()
202 return std::move(body_computation); in MakeCountedLoopBodyComputation()
Dhlo_rematerialization_test.cc238 HloComputation* body_computation = module->AddEmbeddedComputation( in TEST_F() local
242 while_cond, /*while_body=*/body_computation)); in TEST_F()
245 EXPECT_EQ(body_computation->instruction_count(), 8); in TEST_F()
257 EXPECT_EQ(body_computation->instruction_count(), 8); in TEST_F()
274 HloComputation* body_computation = module->AddEmbeddedComputation( in TEST_F() local
278 while_cond, /*while_body=*/body_computation)); in TEST_F()
281 EXPECT_EQ(body_computation->instruction_count(), 8); in TEST_F()
290 EXPECT_EQ(body_computation->instruction_count(), 9); in TEST_F()
Dhlo_memory_scheduler_test.cc272 auto body_computation = module->AddEmbeddedComputation(body_builder.Build()); in TEST_F() local
280 r1f32, cond_computation, body_computation, while_init)); in TEST_F()
297 memory_by_computation[body_computation] = 16; in TEST_F()
Dheap_simulator_test.cc66 HloComputation* body_computation = in TEST_F() local
81 tuple_shape, cond_computation, body_computation, tuple)); in TEST_F()
92 schedule.set_sequence(body_computation, {body_param}); in TEST_F()
163 auto body_computation = module->AddEmbeddedComputation(body_builder.Build()); in TEST_F() local
173 r1f32, cond_computation, body_computation, while_init)); in TEST_F()
200 schedule.set_sequence(body_computation, while_body_vec); in TEST_F()
208 memory_by_computation[body_computation] = 16; in TEST_F()
715 HloComputation* body_computation = in TEST_F() local
722 tuple_shape, cond_computation, body_computation, param)); in TEST_F()
Dhlo_dce_test.cc243 auto body_computation = module->AddEmbeddedComputation(body_builder.Build()); in TEST_F() local
250 shape, cond_computation, body_computation, param)); in TEST_F()
Dbuffer_assignment_test.cc853 auto body_computation = in TEST_F() local
865 t_s32_f32v4_, condition_computation, body_computation, tuple)); in TEST_F()
874 GetInstructions(body_computation->root_instruction()); in TEST_F()
895 HloInstruction* body_root = body_computation->root_instruction(); in TEST_F()
Dtuple_points_to_analysis_test.cc1205 HloComputation* body_computation = in TEST_F() local
1212 data_shape, cond_computation, body_computation, data)); in TEST_F()
Dhlo_dataflow_analysis_test.cc2611 HloComputation* body_computation = in TEST_F() local
2618 data_shape, cond_computation, body_computation, data)); in TEST_F()
/external/tensorflow/tensorflow/compiler/xla/tests/
Dwhile_test.cc969 TF_ASSERT_OK_AND_ASSIGN(auto body_computation, body.Build()); in XLA_TEST_F()
970 While(cond_computation, body_computation, t); in XLA_TEST_F()
997 TF_ASSERT_OK_AND_ASSIGN(auto body_computation, body.Build()); in XLA_TEST_F()
998 While(cond_computation, body_computation, p); in XLA_TEST_F()
1023 TF_ASSERT_OK_AND_ASSIGN(auto body_computation, body.Build()); in XLA_TEST_F()
1024 While(cond_computation, body_computation, p); in XLA_TEST_F()
1062 TF_ASSERT_OK_AND_ASSIGN(auto body_computation, body.Build()); in XLA_TEST_F()
1063 While(cond_computation, body_computation, p); in XLA_TEST_F()