Searched refs:inv_diag_blocks (Results 1 – 1 of 1) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | triangular_solve_expander.cc | 241 auto inv_diag_blocks = GetTupleElement(invert_while, 1); in InvertDiagonalBlocks() local
244 inv_diag_blocks = Div(inv_diag_blocks, diags, in InvertDiagonalBlocks()
248 return Reshape(inv_diag_blocks, AsInt64Slice(shape.dimensions())); in InvertDiagonalBlocks()
252 XlaOp SolveWithInvertedDiagonalBlocks(XlaOp a, XlaOp b, XlaOp inv_diag_blocks, in SolveWithInvertedDiagonalBlocks() argument
258 TF_ASSIGN_OR_RETURN(Shape blocks_shape, builder->GetShape(inv_diag_blocks)); in SolveWithInvertedDiagonalBlocks()
288 MaybeConjugate(Collapse(SliceInMinorDims(inv_diag_blocks, {j, 0, 0}, in SolveWithInvertedDiagonalBlocks()
433 auto inv_diag_blocks = InvertDiagonalBlocks(diag_blocks, lower, transpose_a, in BuildTriangularSolve() local
438 SolveWithInvertedDiagonalBlocks(a, b, inv_diag_blocks, left_side, lower, in BuildTriangularSolve()
|