1# Copyright 2015 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"""Tests for tensorflow.ops.tf.Cholesky.""" 16 17from __future__ import absolute_import 18from __future__ import division 19from __future__ import print_function 20 21import numpy as np 22from six.moves import xrange # pylint: disable=redefined-builtin 23 24from tensorflow.python.client import session 25from tensorflow.python.framework import constant_op 26from tensorflow.python.framework import dtypes as dtypes_lib 27from tensorflow.python.framework import errors_impl 28from tensorflow.python.framework import ops 29from tensorflow.python.framework import test_util 30from tensorflow.python.ops import array_ops 31from tensorflow.python.ops import control_flow_ops 32from tensorflow.python.ops import gradient_checker_v2 33from tensorflow.python.ops import linalg_ops 34from tensorflow.python.ops import math_ops 35from tensorflow.python.ops import stateless_random_ops 36from tensorflow.python.ops import variables 37from tensorflow.python.ops.linalg import linalg 38from tensorflow.python.platform import benchmark 39from tensorflow.python.platform import test 40 41 42# Different gradient implementations for benchmark purposes 43def _GradWithInverseL(l, l_inverse, grad): 44 middle = math_ops.matmul(l, grad, adjoint_a=True) 45 middle = array_ops.matrix_set_diag(middle, 46 0.5 * array_ops.matrix_diag_part(middle)) 47 middle = array_ops.matrix_band_part(middle, -1, 0) 48 grad_a = math_ops.matmul( 49 math_ops.matmul(l_inverse, middle, adjoint_a=True), l_inverse) 50 grad_a += math_ops.conj(array_ops.matrix_transpose(grad_a)) 51 return grad_a * 0.5 52 53 54def TriAngSolveCompositeGrad(l, grad): 55 # Gradient is l^{-H} @ ((l^{H} @ grad) * (tril(ones)-1/2*eye)) @ l^{-1} 56 57 # Compute ((l^{H} @ grad) * (tril(ones)-1/2*eye)) = middle 58 middle = math_ops.matmul(l, grad, adjoint_a=True) 59 middle = array_ops.matrix_set_diag(middle, 60 0.5 * array_ops.matrix_diag_part(middle)) 61 middle = array_ops.matrix_band_part(middle, -1, 0) 62 63 # Compute l^{-H} @ middle = z 64 l_inverse_middle = linalg_ops.matrix_triangular_solve(l, middle, adjoint=True) 65 66 # We need to compute z @ l^{-1}. With matrix_triangular_solve we 67 # actually compute l^{-H} @ z^{H} = grad. Since we later add grad^{H} 68 # we can ommit the conjugate transpose here. 69 z_h = math_ops.conj(array_ops.matrix_transpose(l_inverse_middle)) 70 grad_a = linalg_ops.matrix_triangular_solve(l, z_h, adjoint=True) 71 grad_a += linalg.adjoint(grad_a) 72 return grad_a * 0.5 73 74 75def MatrixInverseCompositeGrad(l, grad): 76 l_inverse = linalg_ops.matrix_inverse(l) 77 return _GradWithInverseL(l, l_inverse, grad) 78 79 80def TriAngInvCompositeGrad(l, grad): 81 num_rows = array_ops.shape(l)[-1] 82 batch_shape = array_ops.shape(l)[:-2] 83 l_inverse = linalg_ops.matrix_triangular_solve(l, 84 linalg_ops.eye( 85 num_rows, 86 batch_shape=batch_shape, 87 dtype=l.dtype)) 88 return _GradWithInverseL(l, l_inverse, grad) 89 90 91class CholeskyOpTest(test.TestCase): 92 93 def _verifyCholeskyBase(self, x, chol, verification): 94 chol_np, verification_np = self.evaluate([chol, verification]) 95 self.assertAllClose(x, verification_np) 96 self.assertShapeEqual(x, chol) 97 # Check that the cholesky is lower triangular, and has positive diagonal 98 # elements. 99 if chol_np.shape[-1] > 0: 100 chol_reshaped = np.reshape(chol_np, (-1, chol_np.shape[-2], 101 chol_np.shape[-1])) 102 for chol_matrix in chol_reshaped: 103 self.assertAllClose(chol_matrix, np.tril(chol_matrix)) 104 self.assertTrue((np.diag(chol_matrix) > 0.0).all()) 105 106 def _verifyCholesky(self, x): 107 # Verify that LL^T == x. 108 chol = linalg_ops.cholesky(x) 109 verification = test_util.matmul_without_tf32(chol, chol, adjoint_b=True) 110 self._verifyCholeskyBase(x, chol, verification) 111 112 @test_util.run_in_graph_and_eager_modes(use_gpu=True) 113 def testBasic(self): 114 data = np.array([[4., -1., 2.], [-1., 6., 0], [2., 0., 5.]]) 115 for dtype in (np.float32, np.float64): 116 with self.subTest(dtype=dtype): 117 self._verifyCholesky(data.astype(dtype)) 118 for dtype in (np.complex64, np.complex128): 119 with self.subTest(dtype=dtype): 120 complex_data = np.tril(1j * data, -1).astype(dtype) 121 complex_data += np.triu(-1j * data, 1).astype(dtype) 122 complex_data += data 123 self._verifyCholesky(complex_data) 124 125 @test_util.run_in_graph_and_eager_modes(use_gpu=True) 126 def testBatch(self): 127 simple_array = np.array([[[1., 0.], [0., 5.]]]) # shape (1, 2, 2) 128 self._verifyCholesky(simple_array) 129 self._verifyCholesky(np.vstack((simple_array, simple_array))) 130 odd_sized_array = np.array([[[4., -1., 2.], [-1., 6., 0], [2., 0., 5.]]]) 131 self._verifyCholesky(np.vstack((odd_sized_array, odd_sized_array))) 132 133 # Generate random positive-definite matrices. 134 matrices = np.random.rand(10, 5, 5) 135 for i in xrange(10): 136 with self.subTest(i=i): 137 matrices[i] = np.dot(matrices[i].T, matrices[i]) 138 self._verifyCholesky(matrices) 139 140 # Generate random complex valued positive-definite matrices. 141 matrices = np.random.rand(10, 5, 5) + 1j * np.random.rand(10, 5, 5) 142 for i in xrange(10): 143 with self.subTest(i=i): 144 matrices[i] = np.dot(matrices[i].T.conj(), matrices[i]) 145 self._verifyCholesky(matrices) 146 147 @test_util.run_in_graph_and_eager_modes(use_gpu=True) 148 def testNonSquareMatrix(self): 149 with self.assertRaises((ValueError, errors_impl.InvalidArgumentError)): 150 linalg_ops.cholesky(np.array([[1., 2., 3.], [3., 4., 5.]])) 151 with self.assertRaises((ValueError, errors_impl.InvalidArgumentError)): 152 linalg_ops.cholesky( 153 np.array([[[1., 2., 3.], [3., 4., 5.]], [[1., 2., 3.], [3., 4., 5.]] 154 ])) 155 156 @test_util.run_in_graph_and_eager_modes(use_gpu=True) 157 def testWrongDimensions(self): 158 tensor3 = constant_op.constant([1., 2.]) 159 with self.assertRaises((ValueError, errors_impl.InvalidArgumentError)): 160 linalg_ops.cholesky(tensor3) 161 with self.assertRaises((ValueError, errors_impl.InvalidArgumentError)): 162 linalg_ops.cholesky(tensor3) 163 164 @test_util.run_in_graph_and_eager_modes(use_gpu=True) 165 def testNotInvertibleCpu(self): 166 # Non-invertible inputs result in lower-triangular NaNs. 167 x = constant_op.constant([[1., -1., 0.], [-1., 1., -1.], [0., -1., 1.]]) 168 chol = linalg_ops.cholesky(x) 169 # Extract the lower-triangular elements. 170 lower_mask = array_ops.matrix_band_part( 171 constant_op.constant(True, shape=x.shape), -1, 0) 172 chol_lower = array_ops.boolean_mask(chol, lower_mask) 173 # Assert all NaN. 174 all_nan = self.evaluate( 175 math_ops.reduce_all(math_ops.reduce_all(math_ops.is_nan(chol_lower)))) 176 self.assertTrue(all_nan) 177 178 @test_util.run_in_graph_and_eager_modes(use_gpu=True) 179 def testEmpty(self): 180 self._verifyCholesky(np.empty([0, 2, 2])) 181 self._verifyCholesky(np.empty([2, 0, 0])) 182 183 @test_util.run_in_graph_and_eager_modes(use_gpu=True) 184 def testConcurrentExecutesWithoutError(self): 185 seed = [42, 24] 186 matrix_shape = [5, 5] 187 matrix1 = stateless_random_ops.stateless_random_normal(matrix_shape, seed) 188 matrix2 = stateless_random_ops.stateless_random_normal(matrix_shape, seed) 189 matrix1 = math_ops.matmul(matrix1, matrix1, adjoint_a=True) 190 matrix2 = math_ops.matmul(matrix2, matrix2, adjoint_a=True) 191 c1 = linalg_ops.cholesky(matrix1) 192 c2 = linalg_ops.cholesky(matrix2) 193 c1_val, c2_val = self.evaluate([c1, c2]) 194 self.assertAllClose(c1_val, c2_val) 195 196 197class CholeskyGradTest(test.TestCase): 198 _backprop_block_size = 16 199 200 def getShapes(self, shapeList): 201 return ((elem, int(np.floor(1.2 * elem))) for elem in shapeList) 202 203 @test_util.run_in_graph_and_eager_modes(use_gpu=True) 204 def testSmallMatrices(self): 205 np.random.seed(0) 206 shapes = self.getShapes([1, 2, 10]) 207 self.runFiniteDifferences( 208 shapes, dtypes=(dtypes_lib.float32, dtypes_lib.float64)) 209 210 @test_util.run_in_graph_and_eager_modes(use_gpu=True) 211 def testSmallMatricesComplex(self): 212 np.random.seed(0) 213 shapes = self.getShapes([1, 2, 10]) 214 self.runFiniteDifferences( 215 shapes, dtypes=(dtypes_lib.complex64, dtypes_lib.complex128)) 216 217 @test_util.run_in_graph_and_eager_modes(use_gpu=True) 218 def testOneBlockMatrices(self): 219 np.random.seed(0) 220 shapes = self.getShapes([self._backprop_block_size + 1]) 221 self.runFiniteDifferences( 222 shapes, 223 dtypes=(dtypes_lib.float32, dtypes_lib.float64), 224 scalar_test=True) 225 226 @test_util.run_in_graph_and_eager_modes(use_gpu=True) 227 def testTwoBlockMatrixFloat(self): 228 np.random.seed(0) 229 shapes = self.getShapes([2 * self._backprop_block_size + 1]) 230 self.runFiniteDifferences( 231 shapes, dtypes=(dtypes_lib.float32,), scalar_test=True) 232 233 @test_util.run_in_graph_and_eager_modes(use_gpu=True) 234 def testTwoBlockMatrixDouble(self): 235 np.random.seed(0) 236 shapes = self.getShapes([2 * self._backprop_block_size + 1]) 237 self.runFiniteDifferences( 238 shapes, dtypes=(dtypes_lib.float64,), scalar_test=True) 239 240 @test_util.run_in_graph_and_eager_modes(use_gpu=True) 241 def testTwoBlockMatrixComplexFloat(self): 242 np.random.seed(0) 243 shapes = self.getShapes([2 * self._backprop_block_size + 1]) 244 self.runFiniteDifferences( 245 shapes, dtypes=(dtypes_lib.complex64,), scalar_test=True) 246 247 @test_util.run_in_graph_and_eager_modes(use_gpu=True) 248 def testTwoBlockMatrixComplexDouble(self): 249 np.random.seed(0) 250 shapes = self.getShapes([2 * self._backprop_block_size + 1]) 251 self.runFiniteDifferences( 252 shapes, dtypes=(dtypes_lib.complex128,), scalar_test=True) 253 254 def _runOneTest(self, shape, dtype, batch, scalar_test): 255 if dtype == dtypes_lib.float64: 256 tol = 1e-5 257 elif dtype == dtypes_lib.complex128: 258 tol = 5e-5 259 else: 260 tol = 5e-3 261 epsilon = np.finfo(dtype.as_numpy_dtype).eps 262 delta = epsilon**(1.0 / 3.0) 263 264 def RandomInput(): 265 a = np.random.randn(shape[0], shape[1]).astype(dtype.as_numpy_dtype) 266 if dtype.is_complex: 267 a += 1j * np.random.randn(shape[0], shape[1]).astype( 268 dtype.as_numpy_dtype) 269 return a 270 271 def Compute(x): 272 # Turn the random matrix x into a Hermitian matrix by 273 # computing the quadratic form x * x^H. 274 a = test_util.matmul_without_tf32( 275 x, math_ops.conj(array_ops.matrix_transpose(x))) / shape[0] 276 if batch: 277 a = array_ops.tile(array_ops.expand_dims(a, 0), [2, 1, 1]) 278 # Finally take the cholesky decomposition of the Hermitian matrix. 279 c = linalg_ops.cholesky(a) 280 if scalar_test: 281 # Reduce to a single scalar output to speed up test. 282 c = math_ops.reduce_mean(c) 283 return c 284 285 theoretical, numerical = gradient_checker_v2.compute_gradient( 286 Compute, [RandomInput()], delta=delta) 287 self.assertAllClose(theoretical, numerical, atol=tol, rtol=tol) 288 289 def runFiniteDifferences(self, 290 shapes, 291 dtypes=(dtypes_lib.float32, dtypes_lib.float64, 292 dtypes_lib.complex64, dtypes_lib.complex128), 293 scalar_test=False): 294 for shape_ in shapes: 295 for dtype_ in dtypes: 296 for batch_ in False, True: 297 self._runOneTest(shape_, dtype_, batch_, scalar_test) 298 299 300class CholeskyBenchmark(test.Benchmark): 301 302 shapes = [ 303 (4, 4), 304 (10, 10), 305 (16, 16), 306 (101, 101), 307 (256, 256), 308 (1000, 1000), 309 (1024, 1024), 310 (2048, 2048), 311 (513, 2, 2), 312 (513, 8, 8), 313 (513, 256, 256), 314 (4, 513, 2, 2), 315 ] 316 317 def _GenerateMatrix(self, shape): 318 batch_shape = shape[:-2] 319 shape = shape[-2:] 320 assert shape[0] == shape[1] 321 n = shape[0] 322 matrix = np.ones(shape).astype(np.float32) / ( 323 2.0 * n) + np.diag(np.ones(n).astype(np.float32)) 324 return np.tile(matrix, batch_shape + (1, 1)) 325 326 def benchmarkCholeskyOp(self): 327 for shape in self.shapes: 328 with ops.Graph().as_default(), \ 329 session.Session(config=benchmark.benchmark_config()) as sess, \ 330 ops.device("/cpu:0"): 331 matrix = variables.Variable(self._GenerateMatrix(shape)) 332 l = linalg_ops.cholesky(matrix) 333 self.evaluate(variables.global_variables_initializer()) 334 self.run_op_benchmark( 335 sess, 336 control_flow_ops.group( 337 l,), 338 min_iters=25, 339 name="cholesky_cpu_{shape}".format(shape=shape)) 340 341 if test.is_gpu_available(True): 342 with ops.Graph().as_default(), \ 343 session.Session(config=benchmark.benchmark_config()) as sess, \ 344 ops.device("/device:GPU:0"): 345 matrix = variables.Variable(self._GenerateMatrix(shape)) 346 l = linalg_ops.cholesky(matrix) 347 self.evaluate(variables.global_variables_initializer()) 348 self.run_op_benchmark( 349 sess, 350 control_flow_ops.group( 351 l,), 352 min_iters=25, 353 name="cholesky_gpu_{shape}".format(shape=shape)) 354 355 def benchmarkGradVariants(self): 356 357 def _BenchmarkGrad(grad_fn, name, device): 358 for shape in self.shapes: 359 matrix = self._GenerateMatrix(shape) 360 with ops.Graph().as_default(), \ 361 session.Session(config=benchmark.benchmark_config()) as sess, \ 362 ops.device(device): 363 l = variables.Variable(np.linalg.cholesky(matrix)) 364 grad_matrix = variables.Variable( 365 np.random.randn(*matrix.shape).astype(np.float32)) 366 grad = grad_fn(l, grad_matrix) 367 self.evaluate(variables.global_variables_initializer()) 368 self.run_op_benchmark( 369 sess, 370 control_flow_ops.group( 371 grad,), 372 min_iters=25, 373 name="{name}_{dev}_{shape}".format( 374 name=name, dev=grad.device, shape=shape)) 375 376 if test.is_gpu_available(True): 377 _BenchmarkGrad(MatrixInverseCompositeGrad, "composite_matrix_inverse", 378 "/device:GPU:0") 379 _BenchmarkGrad(TriAngInvCompositeGrad, "composite_tri_ang_inverse", 380 "/device:GPU:0") 381 _BenchmarkGrad(TriAngSolveCompositeGrad, "composite_triangular_solve", 382 "/device:GPU:0") 383 384 _BenchmarkGrad(MatrixInverseCompositeGrad, "composite_matrix_inverse", 385 "/cpu:0") 386 _BenchmarkGrad(TriAngInvCompositeGrad, "composite_tri_ang_inverse", 387 "/cpu:0") 388 _BenchmarkGrad(TriAngSolveCompositeGrad, "composite_triangular_solve", 389 "/cpu:0") 390 391 392if __name__ == "__main__": 393 test.main() 394