[TF:XLA] Complete the TriangularSolve implementation.
The previous version only handled the case of left_side=false, lower=true, transpose_a=true, conjugate_a=false. This updated implementation handles all 16 combinations of those boolean options, and also instantiates the corresponding MatrixTriangularSolve TF op. To improve compile times and potentially FLOP performance, when lower=true the within-block subroutine used on the diagonal blocks is now a left-looking variant implemented with an XLA HLO While loop. This update also slightly generalizes BatchDot in tf2xla/lib to accept separate conjugation arguments (in addition to transpose arguments). PiperOrigin-RevId: 183756639
Loading
Please sign in to comment