def test_matrix_power(self): np.random.seed(42) # Testing with small values as we want case that would # result in underflow when ordinary `einsum` is used. a = _random_sparse_tensor([200, 4, 4], max_value=1e-20) m = tf.math.log(a) m2 = _logmatmulexp(m, m) m4 = _logmatmulexp(m2, m2) m8 = _logmatmulexp(m4, m4) m16 = _logmatmulexp(m8, m8) u = _logmatmulexp(m16, m16) # Compute product of 32 matrices as a single `logeinsumexp` over # 33 dimensions ensuring that: # (1) a 33-dimensional intermediate isn't materialized # (2) we can work with logs of quantities that would otherwise # cause an underflow. # (3) we can in fact do this with a batch of matrices. letters = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ' lhs = ','.join(['Z' + letters[i] + letters[i + 1] for i in range(32)]) rhs = 'Z' + letters[0] + letters[32] formula = '{}->{}'.format(lhs, rhs) log_as = 32 * [m] v = logeinsumexp.logeinsumexp(formula, *log_as) self.assertAllClose(u, v)
def test_zero_zero_multiplication(self): """Batch matrix multiplication test.""" a = np.array([[0.0, 0.0], [0.0, 0.0]]) v = tf.exp(logeinsumexp('ij,jk->ik', tf.math.log(a), tf.math.log(a))) self.assertAllClose(a, v)
def test_einsum(self, data): formula, tensors = data tensors = [tf.convert_to_tensor(a, dtype=tf.float64) for a in tensors] u = tf.math.log(tf.einsum(formula, *tensors)) log_tensors = [tf.math.log(t) for t in tensors] v = logeinsumexp.logeinsumexp(formula, *log_tensors) self.assertAllClose(u, v)
def test_batch_sum(self): np.random.seed(42) a = _random_sparse_tensor([100, 100]) u = tf.reduce_logsumexp(tf.math.log(a), axis=-1) v = logeinsumexp.logeinsumexp('ij->i', tf.math.log(a)) self.assertAllClose(u, v)
def test_sum(self): np.random.seed(42) a = _random_sparse_tensor([100]) u = tf.reduce_logsumexp(tf.math.log(a)) v = logeinsumexp('i->', tf.math.log(a)) self.assertAllClose(u, v)
def test_permutation(self): """Permutation test.""" np.random.seed(42) t = _random_sparse_tensor(8 * [2]) u = tf.transpose(t, perm=[2, 3, 4, 5, 6, 7, 0, 1]) v = logeinsumexp.logeinsumexp('abcdefgh->cdefghab', t) self.assertAllClose(u, v)
def test_generalized_trace(self): """Batch matrix multiplication test.""" np.random.seed(42) a = _random_sparse_tensor([2, 2, 2, 2, 2, 2, 2, 2]) u = tf.math.log(a[0, 0, 0, 0, 0, 0, 0, 0] + a[1, 1, 1, 1, 1, 1, 1, 1]) v = logeinsumexp.logeinsumexp('iiiiiiii->', tf.math.log(a)) self.assertAllClose(u, v)
def test_batch_diagonal(self): """Batch matrix multiplication test.""" np.random.seed(42) a = _random_sparse_tensor([100, 8, 8]) u = tf.linalg.diag_part(a) v = logeinsumexp.logeinsumexp('ijj->ij', a) self.assertAllClose(u, v)
def test_permutation(self): """Permutation test.""" np.random.seed(42) a = _random_sparse_tensor(8 * [2]) u = tf.transpose(a, perm=[1, 2, 3, 4, 5, 6, 7, 0]) v = logeinsumexp.logeinsumexp('abcdefgh->bcdefgha', a) self.assertAllClose(u, v)
def test_compare_einsum(self): """Batch matrix multiplication test.""" np.random.seed(42) # We're comparing with regular `einsum` so we choose a `min_value` # to make underflow impossible. a = _random_sparse_tensor([2, 2, 2, 2, 2, 2, 2], min_value=0.1) b = _random_sparse_tensor([2, 2, 2, 2, 2, 2, 2], min_value=0.1) formula = 'abcdcfg,edfcbaa->bd' u = tf.math.log(tf.einsum(formula, a, b)) v = logeinsumexp.logeinsumexp(formula, tf.math.log(a), tf.math.log(b)) self.assertAllClose(u, v)
def test_counterexample(self): a = tf.convert_to_tensor(np.array([0., -1000.]), dtype=tf.float64) b = tf.convert_to_tensor(np.array([-1000., 0.]), dtype=tf.float64) v = logeinsumexp.logeinsumexp('i,i->', a, b) self.assertAllClose(np.log(2.) - 1000., v)