def testMultiplyBroadcasting(self):
tt_a = initializers.random_tensor_batch((3, 3, 3), tt_rank=2, batch_size=1,
dtype=self.dtype)
tt_b = initializers.random_tensor_batch((3, 3, 3), tt_rank=2, batch_size=5,
dtype=self.dtype)
with self.test_session() as sess:
res_actual = ops.full(ops.multiply(tt_a, tt_b))
res_actual2 = ops.full(ops.multiply(tt_b, tt_a))
res_desired = ops.full(tt_a) * ops.full(tt_b)
to_run = [res_actual, res_actual2, res_desired]
res_actual_val, res_actual2_val, res_desired_val = sess.run(to_run)
self.assertAllClose(res_actual_val, res_desired_val)
self.assertAllClose(res_actual2_val, res_desired_val)
def testAddBroadcasting(self):
# Sum two TT-tensors with broadcasting.
tt_a = initializers.random_tensor_batch((2, 1, 4), tt_rank=2, batch_size=1,
dtype=self.dtype)
tt_b = initializers.random_tensor_batch((2, 1, 4), tt_rank=[1, 2, 4, 1],
batch_size=3, dtype=self.dtype)
with self.test_session() as sess:
res_actual = ops.full(ops.add(tt_a, tt_b))
res_actual2 = ops.full(tt_b + tt_a)
res_desired = ops.full(tt_a) + ops.full(tt_b)
to_run = [res_actual, res_actual2, res_desired]
res_actual_val, res_actual2_val, res_desired_val = sess.run(to_run)
self.assertAllClose(res_actual_val, res_desired_val)
self.assertAllClose(res_actual2_val, res_desired_val)
def testPairwiseFlatInnerTensor(self):
# Test pairwise_flat_inner of a batch of TT tensors.
tt_tensors_1 = initializers.random_tensor_batch((2, 3, 2),
batch_size=5)
tt_tensors_2 = initializers.random_tensor_batch((2, 3, 2),
batch_size=5)
res_actual = batch_ops.pairwise_flat_inner(tt_tensors_1, tt_tensors_2)
full_tensors_1 = tf.reshape(ops.full(tt_tensors_1), (5, 12))
full_tensors_2 = tf.reshape(ops.full(tt_tensors_2), (5, 12))
res_desired = tf.matmul(full_tensors_1, tf.transpose(full_tensors_2))
res_desired = tf.squeeze(res_desired)
with self.test_session() as sess:
res_actual_val, res_desired_val = sess.run(
(res_actual, res_desired))
self.assertAllClose(res_desired_val, res_actual_val)
def testMultiplyUnknownSizeBatchAndBatch(self):
c1 = tf.placeholder(tf.float32, [None, 1, 3, 2])
c2 = tf.placeholder(tf.float32, [None, 2, 3, 1])
tt_b = initializers.random_tensor_batch((3, 3),
tt_rank=2,
batch_size=8)
tt_a = TensorTrainBatch([c1, c2])
res_ab = ops.full(ops.multiply(tt_a, tt_b))
res_ba = ops.full(ops.multiply(tt_b, tt_a))
res_desired = ops.full(tt_a) * ops.full(tt_b)
to_run = [res_ab, res_ba, res_desired]
feed_dict = {
c1: np.random.rand(8, 1, 3, 2),
c2: np.random.rand(8, 2, 3, 1)
}
feed_dict_err = {
c1: np.random.rand(1, 1, 3, 2),
c2: np.random.rand(1, 2, 3, 1)
}
with self.test_session() as sess:
ab_full, ba_full, des_full = sess.run(to_run, feed_dict=feed_dict)
self.assertAllClose(ab_full, des_full)
self.assertAllClose(ba_full, des_full)
with self.assertRaises(tf.errors.InvalidArgumentError):
sess.run(to_run, feed_dict=feed_dict_err)
def testMultiplyUnknownBatchSizeBroadcasting(self):
c1 = tf.placeholder(tf.float32, [None, 1, 3, 2])
c2 = tf.placeholder(tf.float32, [None, 2, 3, 1])
tt_a = TensorTrainBatch([c1, c2])
tt_b = initializers.random_tensor_batch((3, 3),
tt_rank=3,
batch_size=1)
tt_c = initializers.random_tensor((3, 3), tt_rank=3)
res_ab = ops.full(ops.multiply(tt_a, tt_b))
res_ba = ops.full(ops.multiply(tt_b, tt_a))
res_ac = ops.full(ops.multiply(tt_a, tt_c))
res_ca = ops.full(ops.multiply(tt_c, tt_a))
res_desired_ab = ops.full(tt_a) * ops.full(tt_b)
res_desired_ac = ops.full(tt_a) * ops.full(tt_c)
to_run = [
res_ab, res_ba, res_ac, res_ca, res_desired_ab, res_desired_ac
]
feed_dict = {
c1: np.random.rand(7, 1, 3, 2),
c2: np.random.rand(7, 2, 3, 1)
}
with self.test_session() as sess:
ab, ba, ac, ca, des_ab, des_ac = sess.run(to_run,
feed_dict=feed_dict)
self.assertAllClose(ab, des_ab)
self.assertAllClose(ba, des_ab)
self.assertAllClose(ac, des_ac)
self.assertAllClose(ca, des_ac)
def testReduceSumBatchMultipleWeighted(self):
# Multiple weighted sums of a batch of TT-tensors.
def desired(tt_batch, coef):
res = coef[0] * tt_batch[0]
for i in range(1, tt_batch.batch_size):
res += coef[i] * tt_batch[i]
return res
with self.test_session() as sess:
tt_batch = initializers.random_tensor_batch((4, 3, 5),
tt_rank=2,
batch_size=3)
coef = [[1., 0.1], [0.9, -0.2], [0.3, 0.3]]
coef = np.array(coef).astype(np.float32)
res_actual = ops.full(
approximate.reduce_sum_batch(tt_batch, 6, coef))
res_desired_1 = ops.full(desired(tt_batch, coef[:, 0]))
res_desired_2 = ops.full(desired(tt_batch, coef[:, 1]))
res_desired = tf.stack((res_desired_1, res_desired_2))
res_desired_val, res_actual_val = sess.run(
[res_desired, res_actual])
self.assertAllClose(res_desired_val,
res_actual_val,
atol=1e-5,
rtol=1e-5)
def testPlaceholderTensorIndexing(self):
tens = initializers.random_tensor_batch((3, 3, 4), batch_size=3)
with self.test_session() as sess:
start = tf.placeholder(tf.int32)
end = tf.placeholder(tf.int32)
desired = ops.full(tens)[0:-1]
actual = ops.full(tens[start:end])
desired, actual = sess.run([desired, actual], {start: 0, end: -1})
self.assertAllClose(desired, actual)
desired = ops.full(tens)[0:1]
actual = ops.full(tens[start:end])
desired, actual = sess.run([desired, actual], {start: 0, end: 1})
self.assertAllClose(desired, actual)
desired = ops.full(tens)[1]
actual = ops.full(tens[start])
desired, actual = sess.run([desired, actual], {start: 1})
self.assertAllClose(desired, actual)
desired = ops.full(tens)[1, 1:3, 1, :3]
actual = ops.full(tens[start, start:end, start, :end])
desired, actual = sess.run([desired, actual], {start: 1, end: 3})
self.assertAllClose(desired, actual)
def testOrthogonalizeLeftToRight(self):
shape = (2, 4, 3, 3)
tt_ranks = (1, 5, 2, 17, 1)
updated_tt_ranks = (1, 2, 2, 6, 1)
tens = initializers.random_tensor_batch(shape,
tt_rank=tt_ranks,
batch_size=2)
orthogonal = decompositions.orthogonalize_tt_cores(tens)
with self.test_session() as sess:
tens_val, orthogonal_val = sess.run(
[ops.full(tens), ops.full(orthogonal)])
self.assertAllClose(tens_val, orthogonal_val, atol=1e-5, rtol=1e-5)
dynamic_tt_ranks = shapes.tt_ranks(orthogonal).eval()
self.assertAllEqual(updated_tt_ranks, dynamic_tt_ranks)
# Check that the TT-cores are orthogonal.
for core_idx in range(4 - 1):
core_shape = (updated_tt_ranks[core_idx] * shape[core_idx],
updated_tt_ranks[core_idx + 1])
for i in range(2):
core = tf.reshape(orthogonal.tt_cores[core_idx][i],
core_shape)
should_be_eye = tf.matmul(tf.transpose(core), core)
should_be_eye_val = sess.run(should_be_eye)
self.assertAllClose(np.eye(updated_tt_ranks[core_idx + 1]),
should_be_eye_val)
def testAddSameBatchSize(self):
# Sum two TT-tensors with the same batch size.
tt_a = initializers.random_tensor_batch((2, 1, 4),
tt_rank=2,
batch_size=3)
tt_b = initializers.random_tensor_batch((2, 1, 4),
tt_rank=[1, 2, 4, 1],
batch_size=3)
with self.test_session() as sess:
res_actual = ops.full(ops.add(tt_a, tt_b))
res_actual2 = ops.full(tt_a + tt_b)
res_desired = ops.full(tt_a) + ops.full(tt_b)
to_run = [res_actual, res_actual2, res_desired]
res_actual_val, res_actual2_val, res_desired_val = sess.run(to_run)
self.assertAllClose(res_actual_val, res_desired_val)
self.assertAllClose(res_actual2_val, res_desired_val)
def testWeightedAddNProjectedBatch(self):
# Add several TT-batches from the same tangent space with coefs.
what1 = initializers.random_tensor_batch((2, 3, 4), 4, batch_size=3)
what2 = initializers.random_tensor_batch((2, 3, 4), 1, batch_size=3)
where = initializers.random_tensor((2, 3, 4), 3)
projected1 = riemannian.project(what1, where)
projected2 = riemannian.project(what2, where)
desired_0 = ops.full(1.2 * projected1[0] + -2.0 * projected2[0])
desired_1 = ops.full(1.9 * projected1[1] + 2.0 * projected2[1])
desired_2 = ops.full(0.0 * projected1[2] + 1.0 * projected2[2])
desired = tf.stack((desired_0, desired_1, desired_2), axis=0)
actual = ops.full(riemannian.add_n_projected((projected1, projected2),
coef=[[1.2, 1.9, 0.0],
[-2.0, 2.0, 1.0]]))
with self.test_session() as sess:
desired_val, actual_val = sess.run((desired, actual))
self.assertAllClose(desired_val, actual_val, atol=1e-5, rtol=1e-5)
def testRandomTensorBatch(self):
shapes = [[3, 4], [3, 4], [3, 4], [3, 4], [1, -2], [1.1, 2], [[3, 4]],
[1, 2], [3, 4]]
tt_ranks = [-2, 1.5, [2, 3, 4, 5], [1.5], 2, 2, 2, 2, 2]
bs = [1] * 7 + [-1] + [0.5]
bad_cases = zip(shapes, tt_ranks, bs)
for case in bad_cases:
with self.assertRaises(ValueError):
initializers.random_tensor_batch(case[0],
tt_rank=case[1],
batch_size=case[2])
for case in bad_cases:
with self.assertRaises(ValueError):
initializers.tensor_batch_with_random_cores(case[0],
tt_rank=case[1],
batch_size=case[2])
with self.assertRaises(NotImplementedError):
initializers.random_tensor_batch([1, 2, 3], mean=1.0)
def testFlatInnerTTTensbyTTTensBroadcasting(self):
# Inner product between two batch TT-tensors with broadcasting.
tt_1 = initializers.random_tensor_batch((2, 3, 4), batch_size=1)
tt_2 = initializers.random_tensor_batch((2, 3, 4), batch_size=3)
res_actual_1 = ops.flat_inner(tt_1, tt_2)
res_actual_2 = ops.flat_inner(tt_2, tt_1)
res_desired = tf.einsum('ijk,oijk->o', ops.full(tt_1[0]),
ops.full(tt_2))
with self.test_session() as sess:
res = sess.run([res_actual_1, res_actual_2, res_desired])
res_actual_1_val, res_actual_2_val, res_desired_val = res
self.assertAllClose(res_actual_1_val, res_desired_val)
self.assertAllClose(res_actual_2_val, res_desired_val)
tt_1 = initializers.random_tensor_batch((2, 3, 4), batch_size=2)
with self.assertRaises(ValueError):
# The batch_sizes are different.
ops.flat_inner(tt_1, tt_2)
def testProjectWeightedSumDtypeBug(self):
# Test that project_sum(TensorTrain, TensorTrain variable, np.array) works.
what = initializers.random_tensor_batch((2, 3, 4), batch_size=3,
dtype=self.dtype)
where = variables.get_variable('a', initializer=what[0])
weights = tf.zeros((3,), dtype=self.dtype)
# Check that it doesn't throw an exception trying to convert weights to
# Variable dtype (float32_ref).
riemannian.project_sum(what, where, weights)
def testFrobeniusNormDifferentiableBatch(self):
with self.test_session() as sess:
tt = initializers.random_tensor_batch((3, 3, 3), tt_rank=2, batch_size=5,
dtype=self.dtype)
norm_sq_diff = ops.frobenius_norm_squared(tt, differentiable=True)
variables = [norm_sq_diff, ops.full(tt)]
norm_sq_diff_val, tt_full = sess.run(variables)
desired_norm = np.linalg.norm(tt_full.reshape((5, -1)), axis=1)**2
self.assertAllClose(norm_sq_diff_val, desired_norm, atol=1e-5, rtol=1e-5)
def testCompareProjectSumAndProject(self):
# Compare results of project_sum and project.
tens = initializers.random_tensor_batch((2, 3, 4), 3, batch_size=4)
tangent_tens = initializers.random_tensor((2, 3, 4), 4)
project_sum = riemannian.project_sum(tens, tangent_tens, tf.eye(4))
project = riemannian.project(tens, tangent_tens)
with self.test_session() as sess:
res = sess.run((ops.full(project_sum), ops.full(project)))
project_sum_val, project_val = res
self.assertAllClose(project_sum_val, project_val)
def testCompareProjectSumAndProject(self):
# Compare results of project_sum and project.
tens = initializers.random_tensor_batch((2, 3, 4), 3, batch_size=4,
dtype=self.dtype)
tangent_tens = initializers.random_tensor((2, 3, 4), 4,
dtype=self.dtype)
project_sum = riemannian.project_sum(tens, tangent_tens, np.eye(4))
project = riemannian.project(tens, tangent_tens)
res = self.evaluate((ops.full(project_sum), ops.full(project)))
project_sum_val, project_val = res
self.assertAllClose(project_sum_val, project_val)
def testConcatTensorPlaceholders(self):
# Test concating TTTensors of unknown batch sizes along batch dimension.
number_of_objects = tf.placeholder(tf.int32)
all = initializers.random_tensor_batch((2, 3), batch_size=5)
actual = batch_ops.concat_along_batch_dim(
(all[:number_of_objects], all[number_of_objects:]))
with self.test_session() as sess:
desired_val, actual_val = sess.run(
(ops.full(all), ops.full(actual)),
feed_dict={number_of_objects: 2})
self.assertAllClose(desired_val, actual_val)
def testGatherNDBatch(self):
idx = [[0, 0, 0, 0], [1, 0, 1, 2], [0, 0, 1, 0]]
tt = initializers.random_tensor_batch((3, 4, 5),
tt_rank=2,
batch_size=2,
dtype=self.dtype)
res_np = ops.gather_nd(tt, idx)
res_desired = tf.gather_nd(ops.full(tt), idx)
to_run = [res_np, res_desired]
res_np_v, des_v = self.evaluate(to_run)
self.assertAllClose(res_np_v, des_v)
def testProjectSum(self):
# Test projecting a batch of TT-tensors.
tens = initializers.random_tensor_batch((2, 3, 4), batch_size=3)
tangent_tens = initializers.random_tensor((2, 3, 4), 3)
weighted_sum = tens[0] + tens[1] + tens[2]
direct_proj = riemannian.project_sum(weighted_sum, tangent_tens)
actual_proj = riemannian.project_sum(tens, tangent_tens)
with self.test_session() as sess:
res = sess.run((ops.full(direct_proj), ops.full(actual_proj)))
desired_val, actual_val = res
self.assertAllClose(desired_val, actual_val)
def testFlatInnerTTTensbyTTTensSameBatchSize(self):
# Inner product between two batch TT-tensors of the same batch_size.
shape_list = ((2, 2),
(2, 3, 4))
rank_list = (1, 2)
with self.test_session() as sess:
for shape in shape_list:
for rank in rank_list:
tt_1 = initializers.random_tensor_batch(shape, tt_rank=rank,
batch_size=2,
dtype=self.dtype)
tt_2 = initializers.random_tensor_batch(shape, tt_rank=rank,
batch_size=2,
dtype=self.dtype)
res_actual = ops.flat_inner(tt_1, tt_2)
tt_1_full = tf.reshape(ops.full(tt_1), (2, 1, -1))
tt_2_full = tf.reshape(ops.full(tt_2), (2, -1, 1))
res_desired = tf.matmul(tt_1_full, tt_2_full)
res_actual_val, res_desired_val = sess.run([res_actual, res_desired])
self.assertAllClose(res_actual_val, np.squeeze(res_desired_val))
def testMultiplyByNumber(self):
# Multiply batch of tensors by a number.
tt = initializers.random_tensor_batch((1, 2, 3), tt_rank=(1, 2, 3, 1),
batch_size=3, dtype=self.dtype)
with self.test_session() as sess:
res_actual = ops.full(ops.multiply(tt, 4))
res_actual2 = ops.full(4.0 * tt)
res_desired = 4.0 * ops.full(tt)
to_run = [res_actual, res_actual2, res_desired]
res_actual_val, res_actual2_val, res_desired_val = sess.run(to_run)
self.assertAllClose(res_actual_val, res_desired_val)
self.assertAllClose(res_actual2_val, res_desired_val)
def testPairwiseFlatInnerTensor(self):
# Compare pairwise_flat_inner_projected against naive implementation.
what1 = initializers.random_tensor_batch((2, 3, 4), 4, batch_size=3)
what2 = initializers.random_tensor_batch((2, 3, 4), 4, batch_size=4)
where = initializers.random_tensor((2, 3, 4), 3)
projected1 = riemannian.project(what1, where)
projected2 = riemannian.project(what2, where)
desired = batch_ops.pairwise_flat_inner(projected1, projected2)
actual = riemannian.pairwise_flat_inner_projected(projected1, projected2)
with self.test_session() as sess:
desired_val, actual_val = sess.run((desired, actual))
self.assertAllClose(desired_val, actual_val, atol=1e-5, rtol=1e-5)
with self.assertRaises(ValueError):
# Second argument is not a projection on the tangent space.
riemannian.pairwise_flat_inner_projected(projected1, what2)
where2 = initializers.random_tensor((2, 3, 4), 3)
another_projected2 = riemannian.project(what2, where2)
with self.assertRaises(ValueError):
# The arguments are projections on different tangent spaces.
riemannian.pairwise_flat_inner_projected(projected1, another_projected2)
def testRoundTensor(self):
shape = (2, 1, 4, 3, 3)
tens = initializers.random_tensor_batch(shape, tt_rank=15, batch_size=3,
dtype=self.dtype)
rounded_tens = decompositions.round(tens, max_tt_rank=9)
vars = [ops.full(tens), ops.full(rounded_tens)]
tens_value, rounded_tens_value = self.evaluate(vars)
# TODO: why so bad accuracy?
self.assertAllClose(tens_value, rounded_tens_value, atol=1e-4,
rtol=1e-4)
dynamic_tt_ranks = self.evaluate(shapes.tt_ranks(rounded_tens))
self.assertAllEqual([1, 2, 2, 8, 3, 1], dynamic_tt_ranks)
def testAddNProjected(self):
# Add several TT-objects from the same tangent space.
what1 = initializers.random_tensor_batch((2, 3, 4), 4, batch_size=3)
what2 = initializers.random_tensor_batch((2, 3, 4), 3, batch_size=3)
where = initializers.random_tensor((2, 3, 4), 3)
projected1 = riemannian.project(what1, where)
projected2 = riemannian.project(what2, where)
desired = ops.full(projected1 + projected2)
actual = ops.full(riemannian.add_n_projected((projected1, projected2)))
with self.test_session() as sess:
desired_val, actual_val = sess.run((desired, actual))
self.assertAllClose(desired_val, actual_val, atol=1e-5, rtol=1e-5)
with self.assertRaises(ValueError):
# Second argument is not a projection on the tangent space.
riemannian.add_n_projected((projected1, what2))
where2 = initializers.random_tensor((2, 3, 4), 3)
another_projected2 = riemannian.project(what2, where2)
with self.assertRaises(ValueError):
# The arguments are projections on different tangent spaces.
riemannian.add_n_projected((projected1, another_projected2))
def testMultiplyBatchByBatch(self):
tt_a = initializers.random_tensor_batch((3, 3, 3),
tt_rank=2,
batch_size=5,
dtype=self.dtype)
tt_b = initializers.random_tensor_batch((3, 3, 3),
tt_rank=2,
batch_size=5,
dtype=self.dtype)
res_actual = ops.full(ops.multiply(tt_a, tt_b))
res_actual2 = ops.full(ops.multiply(tt_b, tt_a))
res_desired = ops.full(tt_a) * ops.full(tt_b)
to_run = [res_actual, res_actual2, res_desired]
res_actual = ops.full(ops.multiply(tt_a, tt_b))
res_actual2 = ops.full(ops.multiply(tt_b, tt_a))
res_desired = ops.full(tt_a) * ops.full(tt_b)
to_run = [res_actual, res_actual2, res_desired]
res_actual_val, res_actual2_val, res_desired_val = self.evaluate(
to_run)
self.assertAllClose(res_actual_val, res_desired_val)
self.assertAllClose(res_actual2_val, res_desired_val)
def testGatherNDBatch(self):
idx = [[0, 0, 0, 0], [1, 0, 1, 2], [0, 0, 1, 0]]
pl_idx = tf.placeholder(tf.int32, [None, 4])
tt = initializers.random_tensor_batch((3, 4, 5), tt_rank=2, batch_size=2,
dtype=self.dtype)
res_np = ops.gather_nd(tt, idx)
res_pl = ops.gather_nd(tt, pl_idx)
res_desired = tf.gather_nd(ops.full(tt), idx)
to_run = [res_np, res_pl, res_desired]
with self.test_session() as sess:
res_np_v, res_pl_v, des_v = sess.run(to_run, feed_dict={pl_idx: idx})
self.assertAllClose(res_np_v, des_v)
self.assertAllClose(res_pl_v, res_pl_v)
def testProjectWeightedSum(self):
# Test projecting a batch of TT-tensors with providing coefs.
tens = initializers.random_tensor_batch((2, 3, 4), 3, batch_size=4)
coef = [0.1, -2, 0, 0.4]
tangent_tens = initializers.random_tensor((2, 3, 4), 4)
weighted_sum = coef[0] * tens[0] + coef[1] * tens[1] + coef[2] * tens[2]
weighted_sum += coef[3] * tens[3]
direct_proj = riemannian.project_sum(weighted_sum, tangent_tens)
actual_proj = riemannian.project_sum(tens, tangent_tens, coef)
with self.test_session() as sess:
res = sess.run((ops.full(direct_proj), ops.full(actual_proj)))
desired_val, actual_val = res
self.assertAllClose(desired_val, actual_val)
def testCoreRenormBatch(self):
a = initializers.random_tensor_batch(3 * (10,), tt_rank=7, batch_size=5,
dtype=self.dtype)
b = ops.renormalize_tt_cores(a)
var_list = [ops.full(a), ops.full(b)]
with self.test_session() as sess:
af, bf = sess.run(var_list)
b_cores = sess.run(b.tt_cores)
b_cores_norms = []
for cr in b_cores:
b_cores_norms.append(np.linalg.norm(cr))
self.assertAllClose(af, bf, atol=1e-5, rtol=1e-5)
self.assertAllClose(b_cores_norms, b_cores_norms[0]
* np.ones((len(b_cores))))
def testFrobeniusNormTens(self):
# Frobenius norm of a batch of TT-tensors.
with self.test_session() as sess:
tt = initializers.random_tensor_batch((2, 1, 3), batch_size=3)
norm_sq_actual = ops.frobenius_norm_squared(tt)
norm_actual = ops.frobenius_norm(tt)
vars = [norm_sq_actual, norm_actual, ops.full(tt)]
norm_sq_actual_val, norm_actual_val, tt_val = sess.run(vars)
tt_val = tt_val.reshape((3, -1))
norm_sq_desired_val = np.sum(tt_val * tt_val, axis=1)
norm_desired_val = np.sqrt(norm_sq_desired_val)
self.assertAllClose(norm_sq_actual_val, norm_sq_desired_val)
self.assertAllClose(norm_actual_val,
norm_desired_val,
atol=1e-5,
rtol=1e-5)
def testReduceSumBatchWeighted(self):
# Weighted sum of a batch of TT-tensors.
def desired(tt_batch, coef):
res = coef[0] * tt_batch[0]
for i in range(1, tt_batch.batch_size):
res += coef[i] * tt_batch[i]
return res
tt_batch = initializers.random_tensor_batch((4, 3, 5),
tt_rank=3,
batch_size=3,
dtype=self.dtype)
res_actual = ops.full(approximate.reduce_sum_batch(tt_batch, 9,
[1.2, -0.2, 1]))
res_desired = ops.full(desired(tt_batch, [1.2, -0.2, 1]))
res_desired_val, res_actual_val = self.evaluate([res_desired, res_actual])
self.assertAllClose(res_desired_val, res_actual_val, atol=1e-5, rtol=1e-5)
