Home
last modified time | relevance | path

Searched refs:body_builder (Results 1 – 18 of 18) 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.cc276 XlaBuilder* body_builder) -> StatusOr<std::vector<XlaOp>> { in HouseHolderBidiagonalization() argument
611 XlaBuilder* body_builder) -> StatusOr<std::vector<XlaOp>> { in WhileLoopFn() argument
696 absl::StrCat(name, "-Inner"), body_builder)); in WhileLoopFn()
Dmath.cc1666 XlaBuilder* body_builder) -> StatusOr<std::vector<XlaOp>> { in LentzThompsonBarnettAlgorithm() argument
1672 body_builder)); in LentzThompsonBarnettAlgorithm()
1678 body_builder)); in LentzThompsonBarnettAlgorithm()
1705 ReduceAll(tolerance_comparison, ConstantR0<bool>(body_builder, false), in LentzThompsonBarnettAlgorithm()
1706 CreateScalarOrComputation(PRED, body_builder)); in LentzThompsonBarnettAlgorithm()
/external/tensorflow/tensorflow/compiler/xla/service/
Dhlo_alias_analysis_test.cc333 auto body_builder = HloComputation::Builder("body"); in TEST_F() local
334 auto body_param = body_builder.AddInstruction( in TEST_F()
336 auto body_element_0 = body_builder.AddInstruction( in TEST_F()
338 auto body_element_1 = body_builder.AddInstruction( in TEST_F()
340 auto add = body_builder.AddInstruction(HloInstruction::CreateBinary( in TEST_F()
342 auto body_tuple = body_builder.AddInstruction( in TEST_F()
344 HloComputation* body = module_->AddEmbeddedComputation(body_builder.Build()); in TEST_F()
519 auto body_builder = HloComputation::Builder("body"); in TEST_F() local
520 auto body_param = body_builder.AddInstruction( in TEST_F()
522 auto body_element_0 = body_builder.AddInstruction( in TEST_F()
[all …]
Dhlo_live_range_test.cc282 auto body_builder = HloComputation::Builder("WhileBody"); in TEST_F() local
283 HloInstruction* body_param = body_builder.AddInstruction( in TEST_F()
285 HloInstruction* body_iter = body_builder.AddInstruction( in TEST_F()
287 HloInstruction* body_data = body_builder.AddInstruction( in TEST_F()
289 HloInstruction* body_iter_increment = body_builder.AddInstruction( in TEST_F()
292 body_builder.AddInstruction(HloInstruction::CreateBinary( in TEST_F()
295 body_builder.AddInstruction(HloInstruction::CreateConstant( in TEST_F()
298 body_builder.AddInstruction(HloInstruction::CreateBinary( in TEST_F()
301 body_builder.AddInstruction(HloInstruction::CreateBinary( in TEST_F()
304 body_builder.AddInstruction(HloInstruction::CreateBinary( in TEST_F()
[all …]
Dcopy_insertion_test.cc1513 auto body_builder = HloComputation::Builder("body"); in TEST_F() local
1514 auto body_param = body_builder.AddInstruction( in TEST_F()
1516 auto body_element_0 = body_builder.AddInstruction( in TEST_F()
1518 auto body_element_1 = body_builder.AddInstruction( in TEST_F()
1520 body_builder.AddInstruction( in TEST_F()
1522 HloComputation* body = module->AddEmbeddedComputation(body_builder.Build()); in TEST_F()
1788 auto body_builder = HloComputation::Builder("body"); in TEST_F() local
1789 auto body_param = body_builder.AddInstruction( in TEST_F()
1791 auto body_element_0 = body_builder.AddInstruction( in TEST_F()
1793 auto body_element_1 = body_builder.AddInstruction( in TEST_F()
[all …]
Dhlo_ordering_test.cc131 auto body_builder = HloComputation::Builder("body"); in TEST_F() local
132 auto body_param = body_builder.AddInstruction( in TEST_F()
134 auto negate = body_builder.AddInstruction(HloInstruction::CreateUnary( in TEST_F()
136 HloComputation* body = module->AddEmbeddedComputation(body_builder.Build()); in TEST_F()
214 auto body_builder = HloComputation::Builder("body"); in TEST_F() local
215 auto body_param = body_builder.AddInstruction( in TEST_F()
217 auto negate = body_builder.AddInstruction(HloInstruction::CreateUnary( in TEST_F()
219 HloComputation* body = module->AddEmbeddedComputation(body_builder.Build()); in TEST_F()
Dhlo_dce_test.cc271 auto body_builder = HloComputation::Builder(TestName() + "-body"); in TEST_F() local
273 auto param = body_builder.AddInstruction( in TEST_F()
275 auto token = body_builder.AddInstruction(HloInstruction::CreateToken()); in TEST_F()
276 auto infeed = body_builder.AddInstruction( in TEST_F()
278 auto infeed_data = body_builder.AddInstruction( in TEST_F()
280 body_builder.AddInstruction(HloInstruction::CreateBinary( in TEST_F()
283 auto body_computation = module->AddEmbeddedComputation(body_builder.Build()); in TEST_F()
Dhlo_dataflow_analysis_test.cc393 auto body_builder = HloComputation::Builder("body"); in TEST_P() local
394 auto body_param = body_builder.AddInstruction( in TEST_P()
396 auto body_element_0 = body_builder.AddInstruction( in TEST_P()
398 auto body_element_1 = body_builder.AddInstruction( in TEST_P()
400 auto add = body_builder.AddInstruction(HloInstruction::CreateBinary( in TEST_P()
402 auto body_root = body_builder.AddInstruction( in TEST_P()
404 HloComputation* body = module_->AddEmbeddedComputation(body_builder.Build()); in TEST_P()
496 auto body_builder = HloComputation::Builder("body"); in TEST_P() local
497 auto body_param = body_builder.AddInstruction( in TEST_P()
499 auto body_element_0 = body_builder.AddInstruction( in TEST_P()
[all …]
Dcholesky_expander.cc79 XlaBuilder* body_builder) -> StatusOr<std::vector<XlaOp>> { in CholeskyUnblocked() argument
86 Iota(body_builder, ShapeUtil::MakeShape(S32, matrix_dims), ndims - 1); in CholeskyUnblocked()
88 Iota(body_builder, ShapeUtil::MakeShape(S32, matrix_dims), ndims - 2); in CholeskyUnblocked()
93 Zeros(body_builder, in CholeskyUnblocked()
Dmemory_space_assignment_test.cc759 auto body_builder = HloComputation::Builder("WhileBody"); in TEST_P() local
761 HloInstruction* body_param = body_builder.AddInstruction( in TEST_P()
763 HloInstruction* body_iter = body_builder.AddInstruction( in TEST_P()
765 HloInstruction* body_data = body_builder.AddInstruction( in TEST_P()
767 HloInstruction* body_iter_increment = body_builder.AddInstruction( in TEST_P()
770 body_builder.AddInstruction(HloInstruction::CreateBinary( in TEST_P()
773 body_builder.AddInstruction(HloInstruction::CreateConstant( in TEST_P()
776 body_builder.AddInstruction(HloInstruction::CreateBinary( in TEST_P()
779 body_builder.AddInstruction(HloInstruction::CreateBinary( in TEST_P()
782 body_builder.AddInstruction(HloInstruction::CreateBinary( in TEST_P()
[all …]
Dheap_simulator_test.cc63 auto body_builder = HloComputation::Builder("WhileBody"); in TEST_F() local
65 HloInstruction* body_param = body_builder.AddInstruction( in TEST_F()
68 module->AddEmbeddedComputation(body_builder.Build()); in TEST_F()
155 auto body_builder = HloComputation::Builder("WhileBody"); in TEST_F() local
156 HloInstruction* body_param = body_builder.AddInstruction( in TEST_F()
159 body_builder.AddInstruction(HloInstruction::CreateConstant( in TEST_F()
162 body_builder.AddInstruction(HloInstruction::CreateBinary( in TEST_F()
164 auto body_computation = module->AddEmbeddedComputation(body_builder.Build()); in TEST_F()
905 auto body_builder = HloComputation::Builder("WhileBody"); in TEST_F() local
906 HloInstruction* body_param = body_builder.AddInstruction( in TEST_F()
[all …]
Ddynamic_dimension_inference_test.cc742 auto body_builder = HloComputation::Builder("body"); in TEST_F() local
743 auto body_param = body_builder.AddInstruction( in TEST_F()
745 auto gte_0 = body_builder.AddInstruction( in TEST_F()
747 auto gte_1 = body_builder.AddInstruction( in TEST_F()
749 auto add = body_builder.AddInstruction( in TEST_F()
751 body_builder.AddInstruction(HloInstruction::CreateTuple({add, add})); in TEST_F()
753 HloComputation* body = module_->AddEmbeddedComputation(body_builder.Build()); in TEST_F()
Dbuffer_assignment_test.cc795 auto body_builder = HloComputation::Builder("WhileBody"); in TEST_F() local
796 HloInstruction* body_param = body_builder.AddInstruction( in TEST_F()
798 HloInstruction* body_iter = body_builder.AddInstruction( in TEST_F()
800 HloInstruction* body_data = body_builder.AddInstruction( in TEST_F()
802 HloInstruction* body_data_increment = body_builder.AddInstruction( in TEST_F()
806 body_builder.AddInstruction(HloInstruction::CreateBinary( in TEST_F()
808 HloInstruction* body_iter_increment = body_builder.AddInstruction( in TEST_F()
811 body_builder.AddInstruction(HloInstruction::CreateBinary( in TEST_F()
813 body_builder.AddInstruction( in TEST_F()
816 module->AddEmbeddedComputation(body_builder.Build()); in TEST_F()
/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/mlir/hlo/lib/Dialect/mhlo/transforms/
Dlhlo_legalize_to_affine.cc35 function_ref<void(OpBuilder&, Location, ValueRange)> body_builder) { in BuildBoundedAffineLoopNest() argument
39 body_builder); in BuildBoundedAffineLoopNest()
75 auto body_builder = [&](OpBuilder& builder, Location loc, ValueRange ivs) { in matchAndRewrite() local
93 body_builder); in matchAndRewrite()
118 auto body_builder = [&](OpBuilder& builder, Location loc, in matchAndRewrite() local
130 body_builder); in matchAndRewrite()
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/translate/
Dimport_model.cc3124 mlir::OpBuilder body_builder(&func.getBody()); in CreateSavedModelIR() local
3125 auto call = body_builder.create<mlir::TF::StatefulPartitionedCallOp>( in CreateSavedModelIR()
3131 body_builder.create<mlir::ReturnOp>(func.getLoc(), call.getResults()); in CreateSavedModelIR()