def _testDenseSparse(self, transpose_a, transpose_b, adjoint_a, adjoint_b): if not self._gpu_available: return sparsify = lambda m: m * (m > 0) dense_shape_a = [5, 13, 7] if transpose_a or adjoint_a else [5, 7, 13] dense_shape_b = [5, 15, 13] if transpose_b or adjoint_b else [5, 13, 15] dtypes_to_test = [np.float32, np.complex64] for dtype in dtypes_to_test: a_mats = (np.random.randn(*dense_shape_a) + 1.j * np.random.randn(*dense_shape_a)).astype(dtype) b_mats = sparsify((np.random.randn(*dense_shape_b) + 1.j * np.random.randn(*dense_shape_b))).astype(dtype) b_sm = sparse_csr_matrix_ops.CSRSparseMatrix(b_mats) c_dense = test_util.matmul_without_tf32( a_mats, b_mats, transpose_a=transpose_a, transpose_b=transpose_b, adjoint_a=adjoint_a, adjoint_b=adjoint_b) c_sm_dense = sparse_csr_matrix_ops.matmul( a_mats, b_sm, transpose_a=transpose_a, transpose_b=transpose_b, adjoint_a=adjoint_a, adjoint_b=adjoint_b) c_dense, c_sm_dense = self.evaluate([c_dense, c_sm_dense]) self.assertAllClose(c_dense, c_sm_dense)
def _verifyCholesky(self, x): # Verify that LL^T == x. chol = linalg_ops.cholesky(x) verification = test_util.matmul_without_tf32(chol, chol, adjoint_b=True) self._verifyCholeskyBase(x, chol, verification)
def _verifyLu(self, x, output_idx_type=dtypes.int64): # Verify that Px = LU. lu, perm = linalg_ops.lu(x, output_idx_type=output_idx_type) # Prepare the lower factor of shape num_rows x num_rows lu_shape = np.array(lu.shape.as_list()) batch_shape = lu_shape[:-2] num_rows = lu_shape[-2] num_cols = lu_shape[-1] lower = array_ops.matrix_band_part(lu, -1, 0) if num_rows > num_cols: eye = linalg_ops.eye(num_rows, batch_shape=batch_shape, dtype=lower.dtype) lower = array_ops.concat([lower, eye[..., num_cols:]], axis=-1) elif num_rows < num_cols: lower = lower[..., :num_rows] # Fill the diagonal with ones. ones_diag = array_ops.ones(np.append(batch_shape, num_rows), dtype=lower.dtype) lower = array_ops.matrix_set_diag(lower, ones_diag) # Prepare the upper factor. upper = array_ops.matrix_band_part(lu, 0, -1) verification = test_util.matmul_without_tf32(lower, upper) # Permute the rows of product of the Cholesky factors. if num_rows > 0: # Reshape the product of the triangular factors and permutation indices # to a single batch dimension. This makes it easy to apply # invert_permutation and gather_nd ops. perm_reshaped = array_ops.reshape(perm, [-1, num_rows]) verification_reshaped = array_ops.reshape(verification, [-1, num_rows, num_cols]) # Invert the permutation in each batch. inv_perm_reshaped = map_fn.map_fn(array_ops.invert_permutation, perm_reshaped) batch_size = perm_reshaped.shape.as_list()[0] # Prepare the batch indices with the same shape as the permutation. # The corresponding batch index is paired with each of the `num_rows` # permutation indices. batch_indices = math_ops.cast(array_ops.broadcast_to( math_ops.range(batch_size)[:, None], perm_reshaped.shape), dtype=output_idx_type) if inv_perm_reshaped.shape == [0]: inv_perm_reshaped = array_ops.zeros_like(batch_indices) permuted_verification_reshaped = array_ops.gather_nd( verification_reshaped, array_ops.stack([batch_indices, inv_perm_reshaped], axis=-1)) # Reshape the verification matrix back to the original shape. verification = array_ops.reshape(permuted_verification_reshaped, lu_shape) self._verifyLuBase(x, lower, upper, perm, verification, output_idx_type)
def _verifySquareRoot(self, matrix, np_type): matrix = matrix.astype(np_type) # Verify that matmul(sqrtm(A), sqrtm(A)) = A sqrt = gen_linalg_ops.matrix_square_root(matrix) square = test_util.matmul_without_tf32(sqrt, sqrt) self.assertShapeEqual(matrix, square) self.assertAllClose(matrix, square, rtol=1e-4, atol=1e-3)
def CheckApproximation(self, a, q, r): if is_single: tol = 1e-5 else: tol = 1e-14 # Tests that a ~= q*r. a_recon = test_util.matmul_without_tf32(q, r) self.assertAllClose(a_recon, a, rtol=tol, atol=tol)
def _verifyCholesky(self, x, atol=1e-6): # Verify that LL^T == x. with self.session() as sess: placeholder = array_ops.placeholder( dtypes.as_dtype(x.dtype), shape=x.shape) with self.test_scope(): chol = linalg_ops.cholesky(placeholder) verification = test_util.matmul_without_tf32(chol, chol, adjoint_b=True) self._verifyCholeskyBase(sess, placeholder, x, chol, verification, atol)
def CheckUnitary(self, x): # Tests that x[...,:,:]^H * x[...,:,:] is close to the identity. xx = test_util.matmul_without_tf32(x, x, adjoint_a=True) identity = array_ops.matrix_band_part(array_ops.ones_like(xx), 0, 0) if is_single: tol = 1e-5 else: tol = 1e-14 self.assertAllClose(identity, xx, atol=tol)
def Test(self): np.random.seed(1) n = shape_[-1] batch_shape = shape_[:-2] np_dtype = dtype_.as_numpy_dtype a = np.random.uniform(low=-1.0, high=1.0, size=n * n).reshape([n, n]).astype(np_dtype) if dtype_.is_complex: a += 1j * np.random.uniform(low=-1.0, high=1.0, size=n * n).reshape([n, n]).astype(np_dtype) a += np.conj(a.T) a = np.tile(a, batch_shape + (1, 1)) if dtype_ in (dtypes_lib.float32, dtypes_lib.complex64): atol = 1e-4 else: atol = 1e-12 np_e, np_v = np.linalg.eigh(a) with self.session(use_gpu=True): if compute_v_: tf_e, tf_v = linalg_ops.self_adjoint_eig( constant_op.constant(a)) # Check that V*diag(E)*V^T is close to A. a_ev = test_util.matmul_without_tf32( test_util.matmul_without_tf32(tf_v, array_ops.matrix_diag(tf_e)), tf_v, adjoint_b=True) self.assertAllClose(self.evaluate(a_ev), a, atol=atol) # Compare to numpy.linalg.eigh. CompareEigenDecompositions(self, np_e, np_v, self.evaluate(tf_e), self.evaluate(tf_v), atol) else: tf_e = linalg_ops.self_adjoint_eigvals(constant_op.constant(a)) self.assertAllClose(np.sort(np_e, -1), np.sort(self.evaluate(tf_e), -1), atol=atol)
def Compute(x): # Turn the random matrix x into a Hermitian matrix by # computing the quadratic form x * x^H. a = test_util.matmul_without_tf32( x, math_ops.conj(array_ops.matrix_transpose(x))) / shape[0] if batch: a = array_ops.tile(array_ops.expand_dims(a, 0), [2, 1, 1]) # Finally take the cholesky decomposition of the Hermitian matrix. c = linalg_ops.cholesky(a) if scalar_test: # Reduce to a single scalar output to speed up test. c = math_ops.reduce_mean(c) return c
def _verifyInverse(self, x, np_type): for adjoint in False, True: y = x.astype(np_type) with self.cached_session(use_gpu=True): # Verify that x^{-1} * x == Identity matrix. inv = linalg_ops.matrix_inverse(y, adjoint=adjoint) tf_ans = test_util.matmul_without_tf32(inv, y, 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 = self.evaluate(tf_ans) self.assertAllClose(np_ans, out, rtol=1e-4, atol=1e-3) self.assertShapeEqual(y, tf_ans)
def _VerifyTriangularSolve(self, a, b, lower, adjoint, atol): clean_a = np.tril(a) if lower else np.triu(a) with self.session() as sess: placeholder_a = MakePlaceholder(a) placeholder_ca = MakePlaceholder(clean_a) placeholder_b = MakePlaceholder(b) with self.test_scope(): x = linalg_ops.matrix_triangular_solve(placeholder_a, placeholder_b, lower=lower, adjoint=adjoint) verification = test_util.matmul_without_tf32(placeholder_ca, x, adjoint_a=adjoint) self._VerifyTriangularSolveBase(sess, placeholder_a, placeholder_ca, placeholder_b, a, clean_a, b, verification, atol)