| 3150 | } |
| 3151 | |
| 3152 | XlaOp TriangularSolve(XlaOp a, XlaOp b, bool left_side, bool lower, |
| 3153 | bool unit_diagonal, |
| 3154 | TriangularSolveOptions::Transpose transpose_a) { |
| 3155 | XlaBuilder* builder = a.builder(); |
| 3156 | return builder->ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 3157 | HloInstructionProto instr; |
| 3158 | TF_ASSIGN_OR_RETURN(const Shape* a_shape, builder->GetShapePtr(a)); |
| 3159 | TF_ASSIGN_OR_RETURN(const Shape* b_shape, builder->GetShapePtr(b)); |
| 3160 | xla::TriangularSolveOptions& options = |
| 3161 | *instr.mutable_triangular_solve_options(); |
| 3162 | options.set_left_side(left_side); |
| 3163 | options.set_lower(lower); |
| 3164 | options.set_unit_diagonal(unit_diagonal); |
| 3165 | options.set_transpose_a(transpose_a); |
| 3166 | TF_ASSIGN_OR_RETURN(Shape shape, ShapeInference::InferTriangularSolveShape( |
| 3167 | *a_shape, *b_shape, options)); |
| 3168 | *instr.mutable_shape() = shape.ToProto(); |
| 3169 | |
| 3170 | return builder->AddInstruction(std::move(instr), |
| 3171 | HloOpcode::kTriangularSolve, {a, b}); |
| 3172 | }); |
| 3173 | } |
| 3174 | |
| 3175 | XlaOp Cholesky(XlaOp a, bool lower) { |
| 3176 | XlaBuilder* builder = a.builder(); |