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)
