# 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.tf.Cholesky.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function import numpy as np from six.moves import xrange # pylint: disable=redefined-builtin from tensorflow.python.client import session from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes as dtypes_lib from tensorflow.python.framework import errors_impl from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import gen_linalg_ops from tensorflow.python.ops import gradient_checker from tensorflow.python.ops import gradients_impl from tensorflow.python.ops import linalg_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import random_ops from tensorflow.python.ops import variables from tensorflow.python.ops.linalg import linalg from tensorflow.python.platform import benchmark from tensorflow.python.platform import test from tensorflow.python.platform import tf_logging # Different gradient implementations for benchmark purposes def SpecializedGrad(l, grad): return gen_linalg_ops.cholesky_grad(l, grad) def _GradWithInverseL(l, l_inverse, grad): middle = math_ops.matmul(l, grad, adjoint_a=True) middle = array_ops.matrix_set_diag(middle, 0.5 * array_ops.matrix_diag_part(middle)) middle = array_ops.matrix_band_part(middle, -1, 0) grad_a = math_ops.matmul( math_ops.matmul(l_inverse, middle, adjoint_a=True), l_inverse) grad_a += math_ops.conj(array_ops.matrix_transpose(grad_a)) return grad_a * 0.5 def TriAngSolveCompositeGrad(l, grad): # Gradient is l^{-H} @ ((l^{H} @ grad) * (tril(ones)-1/2*eye)) @ l^{-1} # Compute ((l^{H} @ grad) * (tril(ones)-1/2*eye)) = middle middle = math_ops.matmul(l, grad, adjoint_a=True) middle = array_ops.matrix_set_diag(middle, 0.5 * array_ops.matrix_diag_part(middle)) middle = array_ops.matrix_band_part(middle, -1, 0) # Compute l^{-H} @ middle = z l_inverse_middle = linalg_ops.matrix_triangular_solve(l, middle, adjoint=True) # We need to compute z @ l^{-1}. With matrix_triangular_solve we # actually compute l^{-H} @ z^{H} = grad. Since we later add grad^{H} # we can ommit the conjugate transpose here. z_h = math_ops.conj(array_ops.matrix_transpose(l_inverse_middle)) grad_a = linalg_ops.matrix_triangular_solve(l, z_h, adjoint=True) grad_a += linalg.adjoint(grad_a) return grad_a * 0.5 def MatrixInverseCompositeGrad(l, grad): l_inverse = linalg_ops.matrix_inverse(l) return _GradWithInverseL(l, l_inverse, grad) def TriAngInvCompositeGrad(l, grad): num_rows = array_ops.shape(l)[-1] batch_shape = array_ops.shape(l)[:-2] l_inverse = linalg_ops.matrix_triangular_solve(l, linalg_ops.eye( num_rows, batch_shape=batch_shape, dtype=l.dtype)) return _GradWithInverseL(l, l_inverse, grad) class CholeskyOpTest(test.TestCase): def _verifyCholeskyBase(self, sess, x, chol, verification): chol_np, verification_np = sess.run([chol, verification]) self.assertAllClose(x, verification_np) self.assertShapeEqual(x, chol) # Check that the cholesky is lower triangular, and has positive diagonal # elements. if chol_np.shape[-1] > 0: chol_reshaped = np.reshape(chol_np, (-1, chol_np.shape[-2], chol_np.shape[-1])) for chol_matrix in chol_reshaped: self.assertAllClose(chol_matrix, np.tril(chol_matrix)) self.assertTrue((np.diag(chol_matrix) > 0.0).all()) def _verifyCholesky(self, x): # Verify that LL^T == x. with self.test_session(use_gpu=True) as sess: chol = linalg_ops.cholesky(x) verification = math_ops.matmul(chol, chol, adjoint_b=True) self._verifyCholeskyBase(sess, x, chol, verification) def testBasic(self): data = np.array([[4., -1., 2.], [-1., 6., 0], [2., 0., 5.]]) for dtype in (np.float32, np.float64): self._verifyCholesky(data.astype(dtype)) for dtype in (np.complex64, np.complex128): complex_data = np.tril(1j * data, -1).astype(dtype) complex_data += np.triu(-1j * data, 1).astype(dtype) complex_data += data self._verifyCholesky(complex_data) def testBatch(self): simple_array = np.array([[[1., 0.], [0., 5.]]]) # shape (1, 2, 2) self._verifyCholesky(simple_array) self._verifyCholesky(np.vstack((simple_array, simple_array))) odd_sized_array = np.array([[[4., -1., 2.], [-1., 6., 0], [2., 0., 5.]]]) self._verifyCholesky(np.vstack((odd_sized_array, odd_sized_array))) # Generate random positive-definite matrices. matrices = np.random.rand(10, 5, 5) for i in xrange(10): matrices[i] = np.dot(matrices[i].T, matrices[i]) self._verifyCholesky(matrices) # Generate random complex valued positive-definite matrices. matrices = np.random.rand(10, 5, 5) + 1j * np.random.rand(10, 5, 5) for i in xrange(10): matrices[i] = np.dot(matrices[i].T.conj(), matrices[i]) self._verifyCholesky(matrices) def testNonSquareMatrix(self): with self.assertRaises(ValueError): linalg_ops.cholesky(np.array([[1., 2., 3.], [3., 4., 5.]])) with self.assertRaises(ValueError): linalg_ops.cholesky( np.array([[[1., 2., 3.], [3., 4., 5.]], [[1., 2., 3.], [3., 4., 5.]] ])) def testWrongDimensions(self): tensor3 = constant_op.constant([1., 2.]) with self.assertRaises(ValueError): linalg_ops.cholesky(tensor3) with self.assertRaises(ValueError): linalg_ops.cholesky(tensor3) def testNotInvertibleCPU(self): # The input should be invertible. with self.test_session(use_gpu=True): with self.assertRaisesRegexp( errors_impl.InvalidArgumentError, "Cholesky decomposition was not successful. The" " input might not be valid."): # All rows of the matrix below add to zero self._verifyCholesky( np.array([[1., -1., 0.], [-1., 1., -1.], [0., -1., 1.]])) def testEmpty(self): self._verifyCholesky(np.empty([0, 2, 2])) self._verifyCholesky(np.empty([2, 0, 0])) def testConcurrentExecutesWithoutError(self): with self.test_session(use_gpu=True) as sess: matrix1 = random_ops.random_normal([5, 5], seed=42) matrix2 = random_ops.random_normal([5, 5], seed=42) matrix1 = math_ops.matmul(matrix1, matrix1, adjoint_a=True) matrix2 = math_ops.matmul(matrix2, matrix2, adjoint_a=True) c1 = linalg_ops.cholesky(matrix1) c2 = linalg_ops.cholesky(matrix2) c1_val, c2_val = sess.run([c1, c2]) self.assertAllEqual(c1_val, c2_val) class CholeskyGradTest(test.TestCase): _backprop_block_size = 32 def getShapes(self, shapeList): return ((elem, int(np.floor(1.2 * elem))) for elem in shapeList) def testSmallMatrices(self): np.random.seed(0) shapes = self.getShapes([1, 2, 10]) self.runFiniteDifferences( shapes, dtypes=(dtypes_lib.float32, dtypes_lib.float64)) def testSmallMatricesComplex(self): np.random.seed(0) shapes = self.getShapes([1, 2, 10]) self.runFiniteDifferences( shapes, dtypes=(dtypes_lib.complex64, dtypes_lib.complex128)) def testOneBlockMatrices(self): np.random.seed(0) shapes = self.getShapes([self._backprop_block_size + 1]) self.runFiniteDifferences( shapes, dtypes=(dtypes_lib.float32, dtypes_lib.float64), scalarTest=True) def testTwoBlockMatrixFloat(self): np.random.seed(0) shapes = self.getShapes([2 * self._backprop_block_size + 1]) self.runFiniteDifferences( shapes, dtypes=(dtypes_lib.float32,), scalarTest=True) def testTwoBlockMatrixDouble(self): np.random.seed(0) shapes = self.getShapes([2 * self._backprop_block_size + 1]) self.runFiniteDifferences( shapes, dtypes=(dtypes_lib.float64,), scalarTest=True) def testTwoBlockMatrixComplexFloat(self): np.random.seed(0) shapes = self.getShapes([2 * self._backprop_block_size + 1]) self.runFiniteDifferences( shapes, dtypes=(dtypes_lib.complex64,), scalarTest=True) def testTwoBlockMatrixComplexDouble(self): np.random.seed(0) shapes = self.getShapes([2 * self._backprop_block_size + 1]) self.runFiniteDifferences( shapes, dtypes=(dtypes_lib.complex128,), scalarTest=True) def testAgainstSpecialized(self): np.random.seed(0) data = np.random.randn(33, 33).astype(np.float32) data = np.matmul(data, data.T) grad_data = np.random.randn(*data.shape).astype(np.float32) with ops.Graph().as_default(), self.test_session(use_gpu=False) as s: x = constant_op.constant(data, dtypes_lib.float32) chol = linalg_ops.cholesky(x) composite_grad = gradients_impl.gradients(chol, x, grad_data)[0] specialized_grad = SpecializedGrad(chol, grad_data) reference, actual = s.run([specialized_grad, composite_grad]) self.assertAllClose(reference, actual) def runFiniteDifferences(self, shapes, dtypes=(dtypes_lib.float32, dtypes_lib.float64, dtypes_lib.complex64, dtypes_lib.complex128), scalarTest=False): with self.test_session(use_gpu=True): for shape in shapes: for batch in False, True: for dtype in dtypes: if not scalarTest: data = np.random.randn(shape[0], shape[1]) if dtype.is_complex: data = data.astype(np.complex64) data += 1j * np.random.randn(shape[0], shape[1]) x = constant_op.constant(data, dtype) tensor = math_ops.matmul( x, math_ops.conj(array_ops.transpose(x))) / shape[0] else: # This is designed to be a faster test for larger matrices. data = np.random.randn() if dtype.is_complex: data = np.complex64(data) data += 1j * np.random.randn() x = constant_op.constant(data, dtype) R = constant_op.constant( np.random.randn(shape[0], shape[1]), dtype) e = math_ops.multiply(R, x) tensor = math_ops.matmul( e, math_ops.conj(array_ops.transpose(e))) / shape[0] # Inner-most matrices in tensor are positive definite. if batch: tensor = array_ops.tile( array_ops.expand_dims(tensor, 0), [4, 1, 1]) y = linalg_ops.cholesky(tensor) if scalarTest: y = math_ops.reduce_mean(y) error = gradient_checker.compute_gradient_error( x, x._shape_as_list(), y, y._shape_as_list()) tf_logging.info("error = %f", error) if dtype == dtypes_lib.float64: self.assertLess(error, 1e-5) elif dtype == dtypes_lib.complex128: self.assertLess(error, 5e-5) else: self.assertLess(error, 5e-3) class CholeskyBenchmark(test.Benchmark): shapes = [ (4, 4), (10, 10), (16, 16), (101, 101), (256, 256), (1000, 1000), (1024, 1024), (2048, 2048), (513, 2, 2), (513, 8, 8), (513, 256, 256), (4, 513, 2, 2), ] def _GenerateMatrix(self, shape): batch_shape = shape[:-2] shape = shape[-2:] assert shape[0] == shape[1] n = shape[0] matrix = np.ones(shape).astype(np.float32) / ( 2.0 * n) + np.diag(np.ones(n).astype(np.float32)) return np.tile(matrix, batch_shape + (1, 1)) def benchmarkCholeskyOp(self): for shape in self.shapes: with ops.Graph().as_default(), \ session.Session(config=benchmark.benchmark_config()) as sess, \ ops.device("/cpu:0"): matrix = variables.Variable(self._GenerateMatrix(shape)) l = linalg_ops.cholesky(matrix) variables.global_variables_initializer().run() self.run_op_benchmark( sess, control_flow_ops.group( l,), min_iters=25, name="cholesky_cpu_{shape}".format(shape=shape)) if test.is_gpu_available(True): with ops.Graph().as_default(), \ session.Session(config=benchmark.benchmark_config()) as sess, \ ops.device("/device:GPU:0"): matrix = variables.Variable(self._GenerateMatrix(shape)) l = linalg_ops.cholesky(matrix) variables.global_variables_initializer().run() self.run_op_benchmark( sess, control_flow_ops.group( l,), min_iters=25, name="cholesky_gpu_{shape}".format(shape=shape)) def benchmarkGradVariants(self): def _BenchmarkGrad(grad_fn, name, device): for shape in self.shapes: matrix = self._GenerateMatrix(shape) with ops.Graph().as_default(), \ session.Session(config=benchmark.benchmark_config()) as sess, \ ops.device(device): l = variables.Variable(np.linalg.cholesky(matrix)) grad_matrix = variables.Variable( np.random.randn(*matrix.shape).astype(np.float32)) grad = grad_fn(l, grad_matrix) variables.global_variables_initializer().run() self.run_op_benchmark( sess, control_flow_ops.group( grad,), min_iters=25, name="{name}_{dev}_{shape}".format( name=name, dev=grad.device, shape=shape)) if test.is_gpu_available(True): _BenchmarkGrad(MatrixInverseCompositeGrad, "composite_matrix_inverse", "/device:GPU:0") _BenchmarkGrad(TriAngInvCompositeGrad, "composite_tri_ang_inverse", "/device:GPU:0") _BenchmarkGrad(TriAngSolveCompositeGrad, "composite_triangular_solve", "/device:GPU:0") _BenchmarkGrad(MatrixInverseCompositeGrad, "composite_matrix_inverse", "/cpu:0") _BenchmarkGrad(TriAngInvCompositeGrad, "composite_tri_ang_inverse", "/cpu:0") _BenchmarkGrad(TriAngSolveCompositeGrad, "composite_triangular_solve", "/cpu:0") _BenchmarkGrad(SpecializedGrad, "specialized", "/cpu:0") if __name__ == "__main__": test.main()