Searched refs:branch_operands (Results 1 – 6 of 6) sorted by relevance
215 absl::Span<const Shape> branch_operands);
2558 absl::Span<const Shape> branch_operands) { in InferConditionalShape() argument2569 TF_RET_CHECK(branch_computations.size() == branch_operands.size()); in InferConditionalShape()2578 branch_operands[j])) { in InferConditionalShape()2581 ShapeUtil::HumanString(branch_operands[j]), in InferConditionalShape()
1904 absl::Span<const XlaOp> branch_operands) { in Conditional() argument1910 std::vector<Shape> branch_operand_shapes(branch_operands.size()); in Conditional()1913 for (int j = 0; j < branch_operands.size(); ++j) { in Conditional()1915 GetShape(branch_operands[j])); in Conditional()1930 for (const XlaOp branch_operand : branch_operands) { in Conditional()3413 absl::Span<const XlaOp> branch_operands) { in Conditional() argument3415 branch_operands); in Conditional()
538 absl::Span<const XlaOp> branch_operands);965 absl::Span<const XlaOp> branch_operands);1808 absl::Span<const XlaOp> branch_operands);
2158 std::vector<BufferAllocation::Slice> branch_operands; in BuildConditionalThunk() local2161 branch_operands.emplace_back(GetAllocationSlice(*hlo->operand(j + 1))); in BuildConditionalThunk()2170 GetAllocationSlice(*hlo->operand(0)), branch_operands, in BuildConditionalThunk()
563 <b> `Conditional(branch_index, branch_computations, branch_operands)` </b>572 | `branch_operands` | sequence of N `XlaOp` | Arguments of type $$ T_0 , |581 will be invoked with `branch_operands[b]` which must be of the same type. The