[TF:XLA] Implement MatrixInverse in terms of QR decomposition and triangular solve

This is rather slow, but allows us to avoid falling back to tensorflow. Once we
have LU decomposition we can switch to using that instead of QR. It also
doesn't support complex numbers, our implementation of QR is not ready for
them.

Nevertheless, I think it's a good start and people can avoid relying on
workarounds for tf.linalg.inv, we can optimize later.

PiperOrigin-RevId: 255058725

tensorflow/compiler/jit/compilability_check_util.cc

@@ -239,8 +239,10 @@ bool RecursiveCompilabilityChecker::OpIsInaccurate(const Node& node) const {
 
 bool RecursiveCompilabilityChecker::OpIsSlow(const Node& node) const {
   // b/128001705: SelfAdjointEigV2 and Svd performance issues.
+  // b/135640736: MatrixInverse performance issues.
   return node.type_string() == "SelfAdjointEigV2" ||
-         node.type_string() == "Svd" || node.type_string() == "Qr";
+         node.type_string() == "Svd" || node.type_string() == "Qr" ||
+         node.type_string() == "MatrixInverse";
 }
 
 bool RecursiveCompilabilityChecker::IsCompilableNode(

tensorflow/compiler/tests/BUILD

@@ -314,6 +314,21 @@ tf_xla_py_test(
     ],
 )
 
+tf_xla_py_test(
+    name = "matrix_inverse_op_test",
+    size = "small",
+    timeout = "moderate",
+    srcs = ["matrix_inverse_op_test.py"],
+    deps = [
+        ":xla_test",
+        "//tensorflow/python:array_ops",
+        "//tensorflow/python:framework",
+        "//tensorflow/python:linalg_ops",
+        "//tensorflow/python:math_ops",
+        "//tensorflow/python:platform_test",
+    ],
+)
+
 tf_xla_py_test(
     name = "matrix_triangular_solve_op_test",
     size = "small",

tensorflow/compiler/tests/matrix_inverse_op_test.py

@@ -0,0 +1,86 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for tensorflow.ops.math_ops.matrix_inverse."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import linalg_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import googletest
+
+
+class InverseOpTest(xla_test.XLATestCase):
+
+  def _verifyInverse(self, x, np_type):
+    for adjoint in False, True:
+      y = x.astype(np_type)
+      with self.session() as sess:
+        # Verify that x^{-1} * x == Identity matrix.
+        p = array_ops.placeholder(dtypes.as_dtype(y.dtype), y.shape, name="x")
+        with self.test_scope():
+          inv = linalg_ops.matrix_inverse(p, adjoint=adjoint)
+          tf_ans = math_ops.matmul(inv, p, adjoint_b=adjoint)
+          np_ans = np.identity(y.shape[-1])
+          if x.ndim > 2:
+            tiling = list(y.shape)
+            tiling[-2:] = [1, 1]
+            np_ans = np.tile(np_ans, tiling)
+        out = sess.run(tf_ans, feed_dict={p: y})
+        self.assertAllClose(np_ans, out, rtol=1e-3, atol=1e-3)
+        self.assertShapeEqual(y, tf_ans)
+
+  def _verifyInverseReal(self, x):
+    for np_type in self.float_types & {np.float64, np.float32}:
+      self._verifyInverse(x, np_type)
+
+  def _makeBatch(self, matrix1, matrix2):
+    matrix_batch = np.concatenate(
+        [np.expand_dims(matrix1, 0),
+         np.expand_dims(matrix2, 0)])
+    matrix_batch = np.tile(matrix_batch, [2, 3, 1, 1])
+    return matrix_batch
+
+  def testNonsymmetric(self):
+    # 2x2 matrices
+    matrix1 = np.array([[1., 2.], [3., 4.]])
+    matrix2 = np.array([[1., 3.], [3., 5.]])
+    self._verifyInverseReal(matrix1)
+    self._verifyInverseReal(matrix2)
+    # A multidimensional batch of 2x2 matrices
+    self._verifyInverseReal(self._makeBatch(matrix1, matrix2))
+
+  def testSymmetricPositiveDefinite(self):
+    # 2x2 matrices
+    matrix1 = np.array([[2., 1.], [1., 2.]])
+    matrix2 = np.array([[3., -1.], [-1., 3.]])
+    self._verifyInverseReal(matrix1)
+    self._verifyInverseReal(matrix2)
+    # A multidimensional batch of 2x2 matrices
+    self._verifyInverseReal(self._makeBatch(matrix1, matrix2))
+
+  def testEmpty(self):
+    self._verifyInverseReal(np.empty([0, 2, 2]))
+    self._verifyInverseReal(np.empty([2, 0, 0]))
+
+
+if __name__ == "__main__":
+  googletest.main()

tensorflow/compiler/tf2xla/kernels/BUILD

@@ -54,6 +54,7 @@ tf_kernel_library(
         "lrn_ops.cc",
         "matmul_op.cc",
         "matrix_band_part_op.cc",
+        "matrix_inverse_op.cc",
         "matrix_set_diag_op.cc",
         "matrix_triangular_solve_op.cc",
         "mirror_pad_op.cc",

tensorflow/compiler/tf2xla/kernels/matrix_inverse_op.cc

@@ -0,0 +1,68 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/xla/client/lib/matrix.h"
+#include "tensorflow/compiler/xla/client/lib/qr.h"
+#include "tensorflow/compiler/xla/client/xla_builder.h"
+
+namespace tensorflow {
+namespace {
+
+class MatrixInverseOp : public XlaOpKernel {
+ public:
+  explicit MatrixInverseOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("adjoint", &adjoint_));
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    const TensorShape input_shape = ctx->InputShape(0);
+    int64 ndims = input_shape.dims();
+    OP_REQUIRES(
+        ctx, ndims >= 2,
+        errors::InvalidArgument("Input must have rank >= 2, got ", ndims));
+    OP_REQUIRES(
+        ctx, input_shape.dim_size(ndims - 2) == input_shape.dim_size(ndims - 1),
+        errors::InvalidArgument("Input matrices must be squares, got",
+                                input_shape.dim_size(ndims - 2),
+                                " != ", input_shape.dim_size(ndims - 1)));
+
+    xla::XlaOp input = xla::MaybeTransposeInMinorDims(ctx->Input(0), adjoint_);
+
+    // TODO(b/111271662): Using LU decomposition instead of QR should be faster.
+    auto qr = xla::QRDecomposition(input, /*full_matrices=*/false);
+    OP_REQUIRES_OK(ctx, qr.status());
+
+    xla::XlaOp output = xla::TriangularSolve(
+        qr.ValueOrDie().r, xla::TransposeInMinorDims(qr.ValueOrDie().q),
+        /*left_side=*/true,
+        /*lower=*/false, /*unit_diagonal=*/false,
+        /*transpose_a=*/
+        xla::TriangularSolveOptions::NO_TRANSPOSE);
+    ctx->SetOutput(0, output);
+  }
+
+ private:
+  bool adjoint_;
+
+  TF_DISALLOW_COPY_AND_ASSIGN(MatrixInverseOp);
+};
+
+// TODO(b/135640736): Allow this for integer and complex types.
+REGISTER_XLA_OP(Name("MatrixInverse").TypeConstraint("T", kFloatTypes),
+                MatrixInverseOp);
+
+}  // namespace
+}  // namespace tensorflow