1 /* Copyright 2018 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 #ifndef TENSORFLOW_COMPILER_XLA_CLIENT_LIB_QR_H_ 17 #define TENSORFLOW_COMPILER_XLA_CLIENT_LIB_QR_H_ 18 19 #include "tensorflow/compiler/xla/client/xla_builder.h" 20 #include "tensorflow/compiler/xla/xla_data.pb.h" 21 22 namespace xla { 23 24 // Computes the QR decompositions of a batch of matrices. That is, 25 // given a (batched) matrix a, computes an orthonormal matrix Q and an 26 // upper-triangular matrix R such that a = QR. 27 // `a` must be a (batched) matrix of size [..., m, n]. 28 struct QrDecomposition { 29 // A matrix with the same shape as the input matrix `a`, whose upper triangle 30 // (inclusive of the diagonal) is the matrix R, and whose lower triangle 31 // (exclusive of the diagonal) contains the elementary Householder reflectors. 32 // This is the same output format as used by LAPACK's xGEQRF routine. 33 XlaOp q_and_r; 34 // A vector of shape [..., min(m, n)] containing the scalar factors of the 35 // elementary Householder reflectors. 36 XlaOp taus; 37 }; 38 39 QrDecomposition Qr(XlaOp a); 40 41 // Given `a` and `taus` as returned by `QRDecomposition`, compute the product of 42 // the elementary Householder reflectors (i.e., the matrix Q of the QR 43 // decomposition). The equivalent LAPACK routine is xORGQR/xUNGQR. 44 XlaOp ProductOfElementaryHouseholderReflectors(XlaOp a, XlaOp taus); 45 46 // Helper that combines `Qr` and `ProductOfElementaryHouseholderReflectors` to 47 // compute explicit matrices `q` and `r`. 48 void QrExplicit(XlaOp a, bool full_matrices, XlaOp& q, XlaOp& r); 49 50 } // namespace xla 51 52 #endif // TENSORFLOW_COMPILER_XLA_CLIENT_LIB_QR_H_ 53