Home
last modified time | relevance | path

Searched refs:branch_computations (Results 1 – 12 of 12) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/
Dshape_inference.cc2557 absl::Span<const ProgramShape> branch_computations, in InferConditionalShape() argument
2565 TF_RET_CHECK(2 == branch_computations.size()); in InferConditionalShape()
2567 TF_RET_CHECK(!branch_computations.empty()); in InferConditionalShape()
2569 TF_RET_CHECK(branch_computations.size() == branch_operands.size()); in InferConditionalShape()
2571 for (int j = 0; j < branch_computations.size(); ++j) { in InferConditionalShape()
2572 if (branch_computations[j].parameters_size() != 1) { in InferConditionalShape()
2575 branch_computations[j].parameters_size()); in InferConditionalShape()
2577 if (!ShapeUtil::Compatible(branch_computations[j].parameters(0), in InferConditionalShape()
2582 ShapeUtil::HumanString(branch_computations[j])); in InferConditionalShape()
2590 if (!ShapeUtil::Compatible(branch_computations[0].result(), in InferConditionalShape()
[all …]
Dshape_inference.h214 absl::Span<const ProgramShape> branch_computations,
Dhlo_instruction.cc971 absl::Span<HloComputation* const> branch_computations, in CreateConditional() argument
976 CHECK_EQ(branch_computations.size(), branch_computation_args.size()); in CreateConditional()
977 for (int i = 0; i < branch_computations.size(); ++i) { in CreateConditional()
978 instruction->called_computations_.push_back(branch_computations[i]); in CreateConditional()
1480 absl::MakeSpan(branch_computations()), in CloneWithNewOperands()
2069 const std::vector<HloComputation*>& HloInstruction::branch_computations() in branch_computations() function in xla::HloInstruction
2300 StrJoin(branch_computations(), ", ", in ExtraAttributesToString()
2345 StrJoin(branch_computations(), ",\n", in ExtraAttributesToString()
Dhlo_instruction.h722 absl::Span<HloComputation* const> branch_computations,
1080 const std::vector<HloComputation*>& branch_computations() const;
Dhlo_parser.cc1451 optional<std::vector<HloComputation*>> branch_computations; in ParseInstructionRhs() local
1465 &branch_computations}; in ParseInstructionRhs()
1471 branch_computations.emplace({*true_computation, *false_computation}); in ParseInstructionRhs()
1473 if (branch_computations->empty() || in ParseInstructionRhs()
1474 operands.size() != branch_computations->size() + 1) { in ParseInstructionRhs()
1479 absl::MakeSpan(*branch_computations), in ParseInstructionRhs()
Dhlo_parser_test.cc982 …onstant, f32[] %constant.1, f32[] %constant.2, f32[] %constant.3), branch_computations={%Negate, %… in CreateTestCases()
1249 …= f32[] conditional(constant, constant.1, constant.2, constant.3), branch_computations={Negate, Id… in CreateTestCases()
Dcopy_insertion.cc323 for (HloComputation* computation : conditional->branch_computations()) { in AddCopiesForConditional()
Dlayout_assignment.cc942 for (auto branch_computation : instruction->branch_computations()) { in CheckLayouts()
/external/tensorflow/tensorflow/compiler/xla/client/
Dxla_builder.cc1903 absl::Span<const XlaComputation* const> branch_computations, in Conditional() argument
1912 branch_computations.size()); in Conditional()
1917 branch_computations[j]->GetProgramShape()); in Conditional()
1925 for (const XlaComputation* branch_computation : branch_computations) { in Conditional()
3412 absl::Span<const XlaComputation* const> branch_computations, in Conditional() argument
3414 return branch_index.builder()->Conditional(branch_index, branch_computations, in Conditional()
Dxla_builder.h537 absl::Span<const XlaComputation* const> branch_computations,
964 absl::Span<const XlaComputation* const> branch_computations,
1807 absl::Span<const XlaComputation* const> branch_computations,
/external/tensorflow/tensorflow/compiler/xla/g3doc/
Doperation_semantics.md563 <b> `Conditional(branch_index, branch_computations, branch_operands)` </b>
569 | `branch_computations` | sequence of N | XlaComputations of type $$ |
575 Executes `branch_computations[branch_index]`, and returns the result. If
578 or >= N, then `branch_computations[N-1]` is executed as the default branch.
580 Each `branch_computations[b]` must take in a single argument of type `T_b` and
582 type of the returned value of each `branch_computations[b]` must be the same.
584 Note that only one of the `branch_computations` will be executed depending on
/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dir_emitter_unnested.cc2089 for (auto branch_computation : conditional->branch_computations()) { in CheckConditionalBuffersShareAllocation()