| 1879 | } |
| 1880 | |
| 1881 | XlaOp XlaBuilder::ConditionalImpl( |
| 1882 | XlaOp branch_index, |
| 1883 | absl::Span<const XlaComputation* const> branch_computations, |
| 1884 | absl::Span<const XlaOp> branch_operands) { |
| 1885 | return ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 1886 | HloInstructionProto instr; |
| 1887 | |
| 1888 | TF_ASSIGN_OR_RETURN(const Shape* branch_index_shape, |
| 1889 | GetShapePtr(branch_index)); |
| 1890 | std::vector<Shape> branch_operand_shapes(branch_operands.size()); |
| 1891 | std::vector<ProgramShape> branch_computation_shapes( |
| 1892 | branch_computations.size()); |
| 1893 | for (int j = 0; j < branch_operands.size(); ++j) { |
| 1894 | TF_ASSIGN_OR_RETURN(branch_operand_shapes[j], |
| 1895 | GetShape(branch_operands[j])); |
| 1896 | TF_ASSIGN_OR_RETURN(branch_computation_shapes[j], |
| 1897 | branch_computations[j]->GetProgramShape()); |
| 1898 | } |
| 1899 | TF_ASSIGN_OR_RETURN(const Shape shape, |
| 1900 | ShapeInference::InferConditionalShape( |
| 1901 | *branch_index_shape, branch_computation_shapes, |
| 1902 | branch_operand_shapes)); |
| 1903 | *instr.mutable_shape() = shape.ToProto(); |
| 1904 | |
| 1905 | for (const XlaComputation* branch_computation : branch_computations) { |
| 1906 | AddCalledComputation(*branch_computation, &instr); |
| 1907 | } |
| 1908 | |
| 1909 | std::vector<XlaOp> operands(1, branch_index); |
| 1910 | for (const XlaOp branch_operand : branch_operands) { |
| 1911 | operands.emplace_back(branch_operand); |
| 1912 | } |
| 1913 | return AddInstruction(std::move(instr), HloOpcode::kConditional, |
| 1914 | absl::MakeSpan(operands)); |
| 1915 | }); |
| 1916 | } |
| 1917 | |
| 1918 | Status XlaBuilder::CheckOpBuilder(XlaOp op) const { |
| 1919 | if (this != op.builder()) { |
nothing calls this directly
no test coverage detected