Home
last modified time | relevance | path

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

/external/tensorflow/tensorflow/compiler/xla/service/
Dwhile_util.cc180 TF_ASSIGN_OR_RETURN(std::unique_ptr<HloComputation> body_computation, in MakeCountedLoopBodyComputation()
183 HloInstruction* one = body_computation->AddInstruction( in MakeCountedLoopBodyComputation()
185 HloInstruction* param = body_computation->parameter_instruction(0); in MakeCountedLoopBodyComputation()
201 body_computation->AddInstruction(HloInstruction::CreateTuple(next_state)); in MakeCountedLoopBodyComputation()
202 body_computation->set_root_instruction(next_state_tuple); in MakeCountedLoopBodyComputation()
204 return std::move(body_computation); in MakeCountedLoopBodyComputation()
Dhlo_rematerialization_test_utils_test.cc71 auto body_computation = MakeRematerializableComputation(); in TEST_F() local
73 while_condition.get(), body_computation.get()); in TEST_F()
Dhlo_rematerialization_test.cc158 HloComputation* body_computation = module->AddEmbeddedComputation( in TEST_F() local
162 while_cond, /*while_body=*/body_computation)); in TEST_F()
165 EXPECT_EQ(body_computation->instruction_count(), 8); in TEST_F()
177 EXPECT_EQ(body_computation->instruction_count(), 8); in TEST_F()
194 HloComputation* body_computation = module->AddEmbeddedComputation( in TEST_F() local
198 while_cond, /*while_body=*/body_computation)); in TEST_F()
201 EXPECT_EQ(body_computation->instruction_count(), 8); in TEST_F()
210 EXPECT_EQ(body_computation->instruction_count(), 9); in TEST_F()
Dhlo_live_range_test.cc308 HloComputation* body_computation = in TEST_F() local
319 tuple_shape, cond_computation, body_computation, tuple)); in TEST_F()
326 schedule.set_sequence(body_computation, in TEST_F()
Dheap_simulator_test.cc67 HloComputation* body_computation = in TEST_F() local
82 tuple_shape, cond_computation, body_computation, tuple)); in TEST_F()
93 schedule.set_sequence(body_computation, {body_param}); in TEST_F()
164 auto body_computation = module->AddEmbeddedComputation(body_builder.Build()); in TEST_F() local
174 r1f32, cond_computation, body_computation, while_init)); in TEST_F()
201 schedule.set_sequence(body_computation, while_body_vec); in TEST_F()
209 memory_by_computation[body_computation] = 16; in TEST_F()
908 HloComputation* body_computation = in TEST_F() local
915 tuple_shape, cond_computation, body_computation, param)); in TEST_F()
Dhlo_dce_test.cc283 auto body_computation = module->AddEmbeddedComputation(body_builder.Build()); in TEST_F() local
290 shape, cond_computation, body_computation, param)); in TEST_F()
Dmemory_space_assignment_test.cc786 HloComputation* body_computation = in TEST_P() local
797 tuple_shape, cond_computation, body_computation, tuple)); in TEST_P()
804 schedule.set_sequence(body_computation, in TEST_P()
2669 HloComputation* body_computation = in TEST_P() local
2682 tuple_shape, cond_computation, body_computation, tuple)); in TEST_P()
2693 schedule.set_sequence(body_computation, in TEST_P()
2980 HloComputation* body_computation = in TEST_P() local
3011 tuple_shape, cond_computation, body_computation, tuple)); in TEST_P()
3022 schedule.set_sequence(body_computation, in TEST_P()
3085 HloComputation* body_computation = in TEST_P() local
[all …]
Dbuffer_assignment_test.cc815 HloComputation* body_computation = in TEST_F() local
828 t_s32_f32v10_color1, cond_computation, body_computation, tuple)); in TEST_F()
1025 auto body_computation = in TEST_F() local
1037 t_s32_f32v4_, condition_computation, body_computation, tuple)); in TEST_F()
1046 GetInstructions(body_computation->root_instruction()); in TEST_F()
1067 HloInstruction* body_root = body_computation->root_instruction(); in TEST_F()
Dhlo_dataflow_analysis_test.cc2744 HloComputation* body_computation = in TEST_F() local
2751 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()