Home
last modified time | relevance | path

Searched refs:body_builder (Results 1 – 14 of 14) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/client/lib/
Dloops.cc62 std::unique_ptr<XlaBuilder> body_builder = in WhileLoopHelper() local
65 auto parameter = Parameter(body_builder.get(), 0, tuple_shape, "parameter"); in WhileLoopHelper()
69 body_function(unpack_tuple(parameter, arity, body_builder.get()), in WhileLoopHelper()
70 body_builder.get())); in WhileLoopHelper()
73 Tuple(body_builder.get(), result); in WhileLoopHelper()
75 TF_ASSIGN_OR_RETURN(auto body, body_builder->Build()); in WhileLoopHelper()
94 XlaBuilder* body_builder) -> StatusOr<std::vector<XlaOp>> { in ForEachIndex() argument
101 ConstantLiteral(body_builder, LiteralUtil::One(num_iterations_type)))); in ForEachIndex()
105 body_function(iteration, values, body_builder)); in ForEachIndex()
Dself_adjoint_eig.cc242 XlaBuilder* body_builder) -> StatusOr<std::vector<XlaOp>> { in WhileLoopFn() argument
321 absl::StrCat(name, "-Inner"), body_builder)); in WhileLoopFn()
Dsvd.cc282 XlaBuilder* body_builder) -> StatusOr<std::vector<XlaOp>> { in HouseHolderBidiagonalization() argument
621 XlaBuilder* body_builder) -> StatusOr<std::vector<XlaOp>> { in WhileLoopFn() argument
706 absl::StrCat(name, "-Inner"), body_builder)); in WhileLoopFn()
/external/tensorflow/tensorflow/compiler/xla/service/
Dhlo_alias_analysis_test.cc339 auto body_builder = HloComputation::Builder("body"); in TEST_F() local
340 auto body_param = body_builder.AddInstruction( in TEST_F()
342 auto body_element_0 = body_builder.AddInstruction( in TEST_F()
344 auto body_element_1 = body_builder.AddInstruction( in TEST_F()
346 auto add = body_builder.AddInstruction(HloInstruction::CreateBinary( in TEST_F()
348 auto body_tuple = body_builder.AddInstruction( in TEST_F()
350 HloComputation* body = module_->AddEmbeddedComputation(body_builder.Build()); in TEST_F()
527 auto body_builder = HloComputation::Builder("body"); in TEST_F() local
528 auto body_param = body_builder.AddInstruction( in TEST_F()
530 auto body_element_0 = body_builder.AddInstruction( in TEST_F()
[all …]
Dcholesky_expander.cc70 XlaBuilder* body_builder) -> StatusOr<std::vector<XlaOp>> { in CholeskyUnblocked() argument
76 Zeros(body_builder, in CholeskyUnblocked()
82 Zeros(body_builder, in CholeskyUnblocked()
86 Iota(body_builder, ShapeUtil::MakeShape(S32, row_shape_dims), in CholeskyUnblocked()
89 Iota(body_builder, ShapeUtil::MakeShape(S32, col_shape_dims), in CholeskyUnblocked()
96 auto zero = ConstantR0<int32>(body_builder, 0); in CholeskyUnblocked()
Dcopy_insertion_test.cc1303 auto body_builder = HloComputation::Builder("body"); in TEST_F() local
1304 auto body_param = body_builder.AddInstruction( in TEST_F()
1306 auto body_element_0 = body_builder.AddInstruction( in TEST_F()
1308 auto body_element_1 = body_builder.AddInstruction( in TEST_F()
1310 body_builder.AddInstruction( in TEST_F()
1312 HloComputation* body = module->AddEmbeddedComputation(body_builder.Build()); in TEST_F()
1585 auto body_builder = HloComputation::Builder("body"); in TEST_F() local
1586 auto body_param = body_builder.AddInstruction( in TEST_F()
1588 auto body_element_0 = body_builder.AddInstruction( in TEST_F()
1590 auto body_element_1 = body_builder.AddInstruction( in TEST_F()
[all …]
Dhlo_ordering_test.cc132 auto body_builder = HloComputation::Builder("body"); in TEST_F() local
133 auto body_param = body_builder.AddInstruction( in TEST_F()
135 auto negate = body_builder.AddInstruction(HloInstruction::CreateUnary( in TEST_F()
137 HloComputation* body = module->AddEmbeddedComputation(body_builder.Build()); in TEST_F()
215 auto body_builder = HloComputation::Builder("body"); in TEST_F() local
216 auto body_param = body_builder.AddInstruction( in TEST_F()
218 auto negate = body_builder.AddInstruction(HloInstruction::CreateUnary( in TEST_F()
220 HloComputation* body = module->AddEmbeddedComputation(body_builder.Build()); in TEST_F()
Dhlo_dataflow_analysis_test.cc486 auto body_builder = HloComputation::Builder("body"); in TEST_P() local
487 auto body_param = body_builder.AddInstruction( in TEST_P()
489 auto body_element_0 = body_builder.AddInstruction( in TEST_P()
491 auto body_element_1 = body_builder.AddInstruction( in TEST_P()
493 auto add = body_builder.AddInstruction(HloInstruction::CreateBinary( in TEST_P()
495 auto body_root = body_builder.AddInstruction( in TEST_P()
497 HloComputation* body = module_->AddEmbeddedComputation(body_builder.Build()); in TEST_P()
589 auto body_builder = HloComputation::Builder("body"); in TEST_P() local
590 auto body_param = body_builder.AddInstruction( in TEST_P()
592 auto body_element_0 = body_builder.AddInstruction( in TEST_P()
[all …]
Dhlo_memory_scheduler_test.cc264 auto body_builder = HloComputation::Builder("WhileBody"); in TEST_F() local
265 HloInstruction* body_param = body_builder.AddInstruction( in TEST_F()
268 body_builder.AddInstruction(HloInstruction::CreateConstant( in TEST_F()
270 body_builder.AddInstruction(HloInstruction::CreateBinary( in TEST_F()
272 auto body_computation = module->AddEmbeddedComputation(body_builder.Build()); in TEST_F()
Dhlo_dce_test.cc233 auto body_builder = HloComputation::Builder(TestName() + "-body"); in TEST_F() local
235 auto param = body_builder.AddInstruction( in TEST_F()
237 auto token = body_builder.AddInstruction(HloInstruction::CreateToken()); in TEST_F()
238 auto infeed = body_builder.AddInstruction( in TEST_F()
240 body_builder.AddInstruction( in TEST_F()
243 auto body_computation = module->AddEmbeddedComputation(body_builder.Build()); in TEST_F()
Ddynamic_dimension_inference_test.cc474 auto body_builder = HloComputation::Builder("body"); in TEST_F() local
475 auto body_param = body_builder.AddInstruction( in TEST_F()
477 auto gte_0 = body_builder.AddInstruction( in TEST_F()
479 auto gte_1 = body_builder.AddInstruction( in TEST_F()
481 auto add = body_builder.AddInstruction( in TEST_F()
483 body_builder.AddInstruction(HloInstruction::CreateTuple({add, add})); in TEST_F()
485 HloComputation* body = module_->AddEmbeddedComputation(body_builder.Build()); in TEST_F()
Dheap_simulator_test.cc62 auto body_builder = HloComputation::Builder("WhileBody"); in TEST_F() local
64 HloInstruction* body_param = body_builder.AddInstruction( in TEST_F()
67 module->AddEmbeddedComputation(body_builder.Build()); in TEST_F()
154 auto body_builder = HloComputation::Builder("WhileBody"); in TEST_F() local
155 HloInstruction* body_param = body_builder.AddInstruction( in TEST_F()
158 body_builder.AddInstruction(HloInstruction::CreateConstant( in TEST_F()
161 body_builder.AddInstruction(HloInstruction::CreateBinary( in TEST_F()
163 auto body_computation = module->AddEmbeddedComputation(body_builder.Build()); in TEST_F()
712 auto body_builder = HloComputation::Builder("WhileBody"); in TEST_F() local
713 HloInstruction* body_param = body_builder.AddInstruction( in TEST_F()
[all …]
/external/tensorflow/tensorflow/compiler/xla/tests/
Doutfeed_in_nested_computation_test.cc52 XlaBuilder body_builder("loop_body"); in XLA_TEST_F() local
53 XlaOp state_tuple = Parameter(&body_builder, 0, state_tuple_shape, "state"); in XLA_TEST_F()
58 XlaOp some_input = Infeed(&body_builder, xfeed_shape); in XLA_TEST_F()
62 XlaOp iter_left = Sub(loop_counter, ConstantR0<int32_t>(&body_builder, 1)); in XLA_TEST_F()
64 Tuple(&body_builder, {iter_left, buffer_inside}); in XLA_TEST_F()
65 return body_builder.Build(); in XLA_TEST_F()
/external/tensorflow/tensorflow/compiler/xla/service/cpu/
Dcpu_copy_insertion_test.cc73 auto body_builder = HloComputation::Builder("body"); in TEST_F() local
74 body_builder.AddInstruction( in TEST_F()
76 body_builder.AddInstruction( in TEST_F()
78 HloComputation* body = module->AddEmbeddedComputation(body_builder.Build()); in TEST_F()