def testMultiplyTwoBatchesUnknownSize(self):
c1 = tf.placeholder(self.dtype, [None, 1, 3, 2])
c2 = tf.placeholder(self.dtype, [None, 2, 3, 1])
c3 = tf.placeholder(self.dtype, [None, 1, 3, 2])
c4 = tf.placeholder(self.dtype, [None, 2, 3, 1])
tt_a = TensorTrainBatch([c1, c2])
tt_b = TensorTrainBatch([c3, c4])
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(7, 1, 3, 2),
c2:np.random.rand(7, 2, 3, 1),
c3:np.random.rand(7, 1, 3, 2),
c4:np.random.rand(7, 2, 3, 1)}
feed_dict_err = {c1:np.random.rand(7, 1, 3, 2),
c2:np.random.rand(7, 2, 3, 1),
c3:np.random.rand(1, 1, 3, 2),
c4: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 multiply(tt_left, right):
"""Returns a TensorTrain corresponding to element-wise product tt_left * right.
The shapes of tt_left and right should coincide.
Args:
tt_left: `TensorTrain`, TT-tensor or TT-matrix
right: `TensorTrain`, TT-tensor or TT-matrix, OR a number.
Returns
a `TensorTrain` object corresponding to the element-wise product of the
arguments.
Raises
ValueError if the arguments shapes do not coincide.
"""
if not isinstance(right, TensorTrainBase):
# Assume right is a number, not TensorTrain.
tt_cores = list(tt_left.tt_cores)
tt_cores[0] = right * tt_cores[0]
out_ranks = tt_left.get_tt_ranks()
else:
ndims = tt_left.ndims()
if tt_left.is_tt_matrix() != right.is_tt_matrix():
raise ValueError('The arguments should be both TT-tensors or both '
'TT-matrices')
if tt_left.get_shape() != right.get_shape():
raise ValueError('The arguments should have the same shape.')
a_ranks = shapes.lazy_tt_ranks(tt_left)
b_ranks = shapes.lazy_tt_ranks(right)
shape = shapes.lazy_raw_shape(tt_left)
is_matrix = tt_left.is_tt_matrix()
tt_cores = []
for core_idx in range(ndims):
a_core = tt_left.tt_cores[core_idx]
b_core = right.tt_cores[core_idx]
left_rank = a_ranks[core_idx] * b_ranks[core_idx]
right_rank = a_ranks[core_idx + 1] * b_ranks[core_idx + 1]
if is_matrix:
curr_core = tf.einsum('aijb,cijd->acijbd', a_core, b_core)
curr_core = tf.reshape(curr_core,
(left_rank, shape[0][core_idx],
shape[1][core_idx], right_rank))
else:
curr_core = tf.einsum('aib,cid->acibd', a_core, b_core)
curr_core = tf.reshape(
curr_core, (left_rank, shape[0][core_idx], right_rank))
tt_cores.append(curr_core)
combined_ranks = zip(tt_left.get_tt_ranks(), right.get_tt_ranks())
out_ranks = [a * b for a, b in combined_ranks]
if isinstance(tt_left, TensorTrain):
return TensorTrain(tt_cores, tt_left.get_raw_shape(), out_ranks)
else:
return TensorTrainBatch(tt_cores, tt_left.get_raw_shape(), out_ranks,
tt_left.batch_size)
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 cast(tt_a, dtype):
"""Casts a tt-tensor to a new type.
Args:
tt_a: `TensorTrain` object.
dtype: The destination type.
Raises:
TypeError: If `tt_a` cannot be cast to the `dtype`.
ValueError: If `tt_a` is not a `TensorTrain` or `TensorTrainBatch`.
"""
res_cores = []
cores = tt_a.tt_cores
for core_idx in range(tt_a.ndims()):
res_cores.append(tf.cast(cores[core_idx], dtype))
res_shape = tt_a.get_raw_shape()
res_ranks = tt_a.get_tt_ranks()
if isinstance(tt_a, TensorTrain):
return TensorTrain(res_cores, res_shape, res_ranks)
elif isinstance(tt_a, TensorTrainBatch):
return TensorTrainBatch(res_cores, res_shape, res_ranks,
tt_a.batch_size)
else:
raise ValueError(
'Unsupported type of input "%s", should be TensorTrain or '
'TensorTrainBatch.' % tt_a)
def multiply_along_batch_dim(batch_tt,
weights,
name='t3f_multiply_along_batch_dim'):
"""Multiply each TensorTrain in a batch by a number.
Args:
batch_tt: TensorTrainBatch object, TT-matrices or TT-tensors.
weights: 1-D tf.Tensor (or something convertible to it like np.array) of size
tt.batch_size with weights.
name: string, name of the Op.
Returns:
TensorTrainBatch
"""
with tf.name_scope(name):
weights = tf.convert_to_tensor(weights, dtype=batch_tt.dtype)
tt_cores = list(batch_tt.tt_cores)
if batch_tt.is_tt_matrix():
weights = weights[:, tf.newaxis, tf.newaxis, tf.newaxis,
tf.newaxis]
else:
weights = weights[:, tf.newaxis, tf.newaxis, tf.newaxis]
tt_cores[0] = weights * tt_cores[0]
out_shape = batch_tt.get_raw_shape()
out_ranks = batch_tt.get_tt_ranks()
out_batch_size = batch_tt.batch_size
return TensorTrainBatch(tt_cores, out_shape, out_ranks, out_batch_size)
def cast(tt, dtype, name='t3f_cast'):
"""Casts a tt-tensor to a new type.
Args:
tt: `TensorTrain` object.
dtype: The destination type.
name: string, name of the Op.
Raises:
TypeError: If `tt` cannot be cast to the `dtype`.
ValueError: If `tt` is not a `TensorTrain` or `TensorTrainBatch`.
"""
with tf.name_scope(name, values=tt.tt_cores):
res_cores = []
cores = tt.tt_cores
for core_idx in range(tt.ndims()):
res_cores.append(tf.cast(cores[core_idx], dtype))
res_shape = tt.get_raw_shape()
res_ranks = tt.get_tt_ranks()
if isinstance(tt, TensorTrain):
return TensorTrain(res_cores, res_shape, res_ranks)
elif isinstance(tt, TensorTrainBatch):
return TensorTrainBatch(res_cores, res_shape, res_ranks, tt.batch_size)
else:
raise ValueError('Unsupported type of input "%s", should be TensorTrain '
'or TensorTrainBatch.' % tt)
def add(tt_a, tt_b):
"""Returns a TensorTrain corresponding to elementwise sum tt_a + tt_b.
The shapes of tt_a and tt_b should coincide.
Supports broadcasting:
add(TensorTrainBatch, TensorTrain)
adds TensorTrain to each element in the batch of TTs in TensorTrainBatch.
Args:
tt_a: `TensorTrain`, `TensorTrainBatch`, TT-tensor, or TT-matrix
tt_b: `TensorTrain`, `TensorTrainBatch`, TT-tensor, or TT-matrix
Returns
a `TensorTrain` object corresponding to the element-wise sum of arguments if
both arguments are `TensorTrain`s.
OR a `TensorTrainBatch` if at least one of the arguments is
`TensorTrainBatch`
Raises
ValueError if the arguments shapes do not coincide
"""
ndims = tt_a.ndims()
if tt_a.is_tt_matrix() != tt_b.is_tt_matrix():
raise ValueError('The arguments should be both TT-tensors or both '
'TT-matrices')
if tt_a.get_raw_shape() != tt_b.get_raw_shape():
raise ValueError('The arguments should have the same shape.')
if not shapes.is_batch_broadcasting_possible(tt_a, tt_b):
raise ValueError(
'The batch sizes are different and not 1, broadcasting is '
'not available.')
is_batch_case = isinstance(tt_a, TensorTrainBatch) or isinstance(
tt_b, TensorTrainBatch)
batch_size = None
if is_batch_case:
if tt_a.is_tt_matrix():
tt_cores, batch_size = _add_batch_matrix_cores(tt_a, tt_b)
else:
tt_cores, batch_size = _add_batch_tensor_cores(tt_a, tt_b)
else:
if tt_a.is_tt_matrix():
tt_cores = _add_matrix_cores(tt_a, tt_b)
else:
tt_cores = _add_tensor_cores(tt_a, tt_b)
out_ranks = [1]
static_a_ranks = tt_a.get_tt_ranks()
static_b_ranks = tt_b.get_tt_ranks()
for core_idx in range(1, ndims):
out_ranks.append(static_a_ranks[core_idx] + static_b_ranks[core_idx])
out_ranks.append(1)
if is_batch_case:
return TensorTrainBatch(tt_cores, tt_a.get_raw_shape(), out_ranks,
batch_size)
else:
return TensorTrain(tt_cores, tt_a.get_raw_shape(), out_ranks)
def matrix_batch_with_random_cores(shape,
tt_rank=2,
batch_size=1,
mean=0.,
stddev=1.):
"""Generate a batch of TT-matrices of given shape with N(mean, stddev^2) cores.
Args:
shape: 2d array, shape[0] is the shape of the matrix row-index,
shape[1] is the shape of the column index.
shape[0] and shape[1] should have the same number of elements (d)
Also supports omitting one of the dimensions for vectors, e.g.
matrix_batch_with_random_cores([[2, 2, 2], None])
and
matrix_batch_with_random_cores([None, [2, 2, 2]])
will create a batch of one 8-element column and row vector correspondingly.
tt_rank: a number or a (d+1)-element array with ranks.
batch_size: an integer.
mean: a number, the mean of the normal distribution used for
initializing TT-cores.
stddev: a number, the standard deviation of the normal distribution used
for initializing TT-cores.
Returns:
TensorTrainBatch containing a batch of TT-matrices of size
np.prod(shape[0]) x np.prod(shape[1])
"""
# TODO: good distribution to init training.
# In case the shape is immutable.
shape = list(shape)
# In case shape represents a vector, e.g. [None, [2, 2, 2]]
if shape[0] is None:
shape[0] = np.ones(len(shape[1]), dtype=int)
# In case shape represents a vector, e.g. [[2, 2, 2], None]
if shape[1] is None:
shape[1] = np.ones(len(shape[0]), dtype=int)
shape = np.array(shape)
tt_rank = np.array(tt_rank)
_validate_input_parameters(is_tensor=False,
shape=shape,
tt_rank=tt_rank,
batch_size=batch_size)
num_dims = shape[0].size
if tt_rank.size == 1:
tt_rank = tt_rank * np.ones(num_dims - 1)
tt_rank = np.concatenate([[1], tt_rank, [1]])
shape = shape.astype(int)
tt_rank = tt_rank.astype(int)
# TODO: variable (name?) scope.
tt_cores = [None] * num_dims
for i in range(num_dims):
curr_core_shape = (batch_size, tt_rank[i], shape[0][i], shape[1][i],
tt_rank[i + 1])
tt_cores[i] = tf.random_normal(curr_core_shape,
mean=mean,
stddev=stddev)
return TensorTrainBatch(tt_cores, shape, tt_rank, batch_size)
def random_matrix_batch(shape, tt_rank=2, batch_size=1):
"""Generate a random batch of TT-matrices of given shape.
Args:
shape: 2d array, shape[0] is the shape of the matrix row-index,
shape[1] is the shape of the column index.
shape[0] and shape[1] should have the same number of elements (d)
Also supports ommiting one of the dimensions for vectors, e.g.
random_matrix_batch([[2, 2, 2], None])
and
random_matrix_batch([None, [2, 2, 2]])
will create a batch of one 8-element column and row vector correspondingly.
tt_rank: a number or a (d+1)-element array with ranks.
batch_size: an integer.
Returns:
TensorTrainBatch containing a batch of TT-matrices of size
np.prod(shape[0]) x np.prod(shape[1])
"""
# TODO: good distribution to init training.
# In case the shape is immutable.
shape = list(shape)
# In case shape represents a vector, e.g. [None, [2, 2, 2]]
if shape[0] is None:
shape[0] = np.ones(len(shape[1]))
# In case shape represents a vector, e.g. [[2, 2, 2], None]
if shape[1] is None:
shape[1] = np.ones(len(shape[0]))
shape = np.array(shape)
tt_rank = np.array(tt_rank)
if len(shape.shape) != 2:
raise ValueError('shape should be 2d array')
if shape[0].size != shape[1].size:
raise ValueError('shape[0] should have the same length as shape[1]')
if np.any(shape.flatten() < 1):
raise ValueError('all elements in `shape` should be positive')
if np.any(tt_rank < 1):
raise ValueError('`rank` should be positive')
if tt_rank.size != 1 and tt_rank.size != (shape[0].size + 1):
raise ValueError('`rank` array has inappropriate size')
if batch_size < 1:
raise ValueError('Batch size should be positive')
num_dims = shape[0].size
if tt_rank.size == 1:
tt_rank = tt_rank * np.ones(num_dims - 1)
tt_rank = np.concatenate([[1], tt_rank, [1]])
# TODO: check that ints?
shape = shape.astype(int)
tt_rank = tt_rank.astype(int)
# TODO: variable (name?) scope.
tt_cores = [None] * num_dims
for i in range(num_dims):
curr_core_shape = (batch_size, tt_rank[i], shape[0][i], shape[1][i],
tt_rank[i + 1])
tt_cores[i] = tf.random_normal(curr_core_shape)
return TensorTrainBatch(tt_cores, shape, tt_rank, batch_size)
def cholesky(kron_a, name='t3f_kronecker_cholesky'):
"""Computes the Cholesky decomposition of a given Kronecker-factorized matrix.
Args:
kron_a: `TensorTrain` or `TensorTrainBatch` object containing a matrix or a
batch of matrices of size N x N, factorized into a Kronecker product of
square matrices (all tt-ranks are 1 and all tt-cores are square). All the
cores must be symmetric positive-definite.
name: string, name of the Op.
Returns:
`TensorTrain` object containing a TT-matrix of size N x N if the argument is
`TensorTrain`
`TensorTrainBatch` object, containing TT-matrices of size N x N if the
argument is `TensorTrainBatch`
Raises:
ValueError if the tt-cores of the provided matrix are not square,
or the tt-ranks are not 1.
"""
if not _is_kron(kron_a):
raise ValueError('The argument should be a Kronecker product '
'(tt-ranks should be 1)')
shapes_defined = kron_a.get_shape().is_fully_defined()
if shapes_defined:
i_shapes = kron_a.get_raw_shape()[0]
j_shapes = kron_a.get_raw_shape()[1]
else:
i_shapes = ops.raw_shape(kron_a)[0]
j_shapes = ops.raw_shape(kron_a)[1]
if shapes_defined:
if i_shapes != j_shapes:
raise ValueError(
'The argument should be a Kronecker product of square '
'matrices (tt-cores must be square)')
is_batch = isinstance(kron_a, TensorTrainBatch)
with tf.name_scope(name):
cho_cores = []
for core_idx in range(kron_a.ndims()):
core = kron_a.tt_cores[core_idx]
if is_batch:
core_cho = tf.linalg.cholesky(core[:, 0, :, :, 0])
core_cho = tf.expand_dims(tf.expand_dims(core_cho, 1), -1)
else:
core_cho = tf.linalg.cholesky(core[0, :, :, 0])
core_cho = tf.expand_dims(tf.expand_dims(core_cho, 0), -1)
cho_cores.append(core_cho)
res_ranks = kron_a.get_tt_ranks()
res_shape = kron_a.get_raw_shape()
if is_batch:
return TensorTrainBatch(cho_cores, res_shape, res_ranks)
else:
return TensorTrain(cho_cores, res_shape, res_ranks)
def testFullTensor2d(self):
np.random.seed(1)
for rank in [1, 2]:
a = np.random.rand(3, 10, rank).astype(self.dtype.as_numpy_dtype)
b = np.random.rand(3, rank, 9).astype(self.dtype.as_numpy_dtype)
tt_cores = (a.reshape(3, 1, 10, rank), b.reshape(3, rank, 9, 1))
desired = np.einsum('oib,obj->oij', a, b)
tf_tens = TensorTrainBatch(tt_cores)
actual = self.evaluate(ops.full(tf_tens))
self.assertAllClose(desired, actual)
def testFullTensor2d(self):
np.random.seed(1)
for rank in [1, 2]:
a = np.random.rand(3, 10, rank).astype(np.float32)
b = np.random.rand(3, rank, 9).astype(np.float32)
tt_cores = (a.reshape(3, 1, 10, rank), b.reshape(3, rank, 9, 1))
desired = np.einsum('oib,obj->oij', a, b)
with self.test_session():
tf_tens = TensorTrainBatch(tt_cores)
actual = ops.full(tf_tens)
self.assertAllClose(desired, actual.eval())
def inv(kron_a):
"""Computes the inverse of a given Kronecker-factorized matrix.
Args:
kron_a: `TensorTrain` or `TensorTrainBatch` object containing a matrix or a
batch of matrices of size N x N, factorized into a Kronecker product of
square matrices (all tt-ranks are 1 and all tt-cores are square).
Returns:
`TensorTrain` object containing a TT-matrix of size N x N if the argument is
`TensorTrain`
`TensorTrainBatch` object, containing TT-matrices of size N x N if the
argument is `TensorTrainBatch`
Raises:
ValueError if the tt-cores of the provided matrix are not square,
or the tt-ranks are not 1.
"""
if not _is_kron(kron_a):
raise ValueError('The argument should be a Kronecker product '
'(tt-ranks should be 1)')
shapes_defined = kron_a.get_shape().is_fully_defined()
if shapes_defined:
i_shapes = kron_a.get_raw_shape()[0]
j_shapes = kron_a.get_raw_shape()[1]
else:
i_shapes = ops.raw_shape(kron_a)[0]
j_shapes = ops.raw_shape(kron_a)[1]
if shapes_defined:
if i_shapes != j_shapes:
raise ValueError(
'The argument should be a Kronecker product of square '
'matrices (tt-cores must be square)')
is_batch = isinstance(kron_a, TensorTrainBatch)
inv_cores = []
for core_idx in range(kron_a.ndims()):
core = kron_a.tt_cores[core_idx]
if is_batch:
core_inv = tf.matrix_inverse(core[:, 0, :, :, 0])
core_inv = tf.expand_dims(tf.expand_dims(core_inv, 1), -1)
else:
core_inv = tf.matrix_inverse(core[0, :, :, 0])
core_inv = tf.expand_dims(tf.expand_dims(core_inv, 0), -1)
inv_cores.append(core_inv)
res_ranks = kron_a.get_tt_ranks()
res_shape = kron_a.get_raw_shape()
if is_batch:
return TensorTrainBatch(inv_cores, res_shape, res_ranks)
else:
return TensorTrain(inv_cores, res_shape, res_ranks)
def renormalize_tt_cores(tt, epsilon=1e-8, name='t3f_renormalize_tt_cores'):
"""Renormalizes TT-cores to make them of the same Frobenius norm.
Doesn't change the tensor represented by `tt` object, but renormalizes the
TT-cores to make further computations more stable.
Args:
tt: `TensorTrain` or `TensorTrainBatch` object
epsilon: parameter for numerical stability of sqrt
name: string, name of the Op.
Returns:
`TensorTrain` or `TensorTrainBatch` which represents the same
tensor as tt, but with all cores having equal norm. In the batch
case applies to each TT in `TensorTrainBatch`.
"""
# TODO: bad way to check if batch or not.
with tf.name_scope(name):
epsilon = tf.convert_to_tensor(epsilon, dtype=tf.float32)
if isinstance(tt, TensorTrain):
new_cores = []
running_log_norm = 0
core_norms = []
for core in tt.tt_cores:
cur_core_norm = tf.sqrt(
tf.maximum(tf.reduce_sum(core**2), epsilon))
core_norms.append(cur_core_norm)
running_log_norm += tf.log(cur_core_norm)
running_log_norm = running_log_norm / tt.ndims()
fact = tf.exp(running_log_norm)
for i, core in enumerate(tt.tt_cores):
new_cores.append(core * fact / core_norms[i])
return TensorTrain(new_cores)
else:
sz = (tt.batch_size, ) + (len(tt.tt_cores[0].shape) - 1) * (1, )
running_core_log_norms = tf.zeros(sz, dtype=tt.dtype)
ax = np.arange(len(tt.tt_cores[0].shape))[1:]
fact_list = []
for core in tt.tt_cores:
cur_core_norm_sq = tf.reduce_sum(core**2,
axis=ax,
keepdims=True)
cur_core_norm = tf.sqrt(tf.maximum(epsilon, cur_core_norm_sq))
fact_list.append(cur_core_norm)
running_core_log_norms += tf.math.log(cur_core_norm)
new_cores = []
exp_fact = tf.exp(running_core_log_norms / tt.ndims())
for i, core in enumerate(tt.tt_cores):
new_cores.append(tf.multiply(core, exp_fact / fact_list[i]))
return TensorTrainBatch(new_cores)
def testFullTensor3d(self):
np.random.seed(1)
for rank_1 in [1, 2]:
a = np.random.rand(3, 10, rank_1).astype(self.dtype.as_numpy_dtype)
b = np.random.rand(3, rank_1, 9, 3).astype(self.dtype.as_numpy_dtype)
c = np.random.rand(3, 3, 8).astype(self.dtype.as_numpy_dtype)
tt_cores = (a.reshape(3, 1, 10, rank_1), b, c.reshape((3, 3, 8, 1)))
# Basically do full by hand.
desired = np.einsum('oia,oajb,obk->oijk', a, b, c)
with self.test_session():
tf_tens = TensorTrainBatch(tt_cores)
actual = ops.full(tf_tens)
self.assertAllClose(desired, actual.eval())
def concat_along_batch_dim(tt_list, name='t3f_concat_along_batch_dim'):
"""Concat all TensorTrainBatch objects along batch dimension.
Args:
tt_list: a list of TensorTrainBatch objects.
name: string, name of the Op.
Returns:
TensorTrainBatch
"""
ndims = tt_list[0].ndims()
if isinstance(tt_list, TensorTrainBase):
# Not a list but just one element, nothing to concat.
return tt_list
for batch_idx in range(len(tt_list)):
if not isinstance(tt_list[batch_idx], TensorTrainBatch):
raise ValueError(
'All objects in the list should be TTBatch objects, got '
'%s' % tt_list[batch_idx])
for batch_idx in range(1, len(tt_list)):
if tt_list[batch_idx].get_raw_shape() != tt_list[0].get_raw_shape():
raise ValueError(
'Shapes of all TT-batch objects should coincide, got %s '
'and %s' % (tt_list[0].get_raw_shape(),
tt_list[batch_idx].get_raw_shape()))
if tt_list[batch_idx].get_tt_ranks() != tt_list[0].get_tt_ranks():
raise ValueError(
'TT-ranks of all TT-batch objects should coincide, got '
'%s and %s' %
(tt_list[0].get_tt_ranks(), tt_list[batch_idx].get_tt_ranks()))
list_of_cores_lists = [tt.tt_cores for tt in tt_list]
all_cores = tuple(itertools.chain.from_iterable(list_of_cores_lists))
with tf.name_scope(name):
res_cores = []
for core_idx in range(ndims):
curr_core = tf.concat([tt.tt_cores[core_idx] for tt in tt_list],
axis=0)
res_cores.append(curr_core)
try:
batch_size = sum([tt.batch_size for tt in tt_list])
except TypeError:
# The batch sizes are not defined and you can't sum Nones.
batch_size = None
return TensorTrainBatch(res_cores, tt_list[0].get_raw_shape(),
tt_list[0].get_tt_ranks(), batch_size)
def tensor_batch_with_random_cores(shape,
tt_rank=2,
batch_size=1,
mean=0.,
stddev=1.,
dtype=tf.float32,
name='t3f_tensor_batch_with_random_cores'):
"""Generate a batch of TT-tensors of given shape with N(mean, stddev^2) cores.
Args:
shape: array representing the shape of the future tensor.
tt_rank: a number or a (d+1)-element array with ranks.
batch_size: an integer.
mean: a number, the mean of the normal distribution used for
initializing TT-cores.
stddev: a number, the standard deviation of the normal distribution used
for initializing TT-cores.
dtype: [tf.float32] dtype of the resulting tensor.
name: string, name of the Op.
Returns:
TensorTrainBatch containing TT-tensors
"""
# TODO: support shape and tt_ranks as TensorShape?.
# TODO: support None as a dimension.
shape = np.array(shape)
tt_rank = np.array(tt_rank)
_validate_input_parameters(is_tensor=True,
shape=shape,
tt_rank=tt_rank,
batch_size=batch_size)
num_dims = shape.size
if tt_rank.size == 1:
tt_rank = tt_rank * np.ones(num_dims - 1, dtype=np.int)
tt_rank = np.insert(tt_rank, 0, 1)
tt_rank = np.append(tt_rank, 1)
tt_rank = tt_rank.astype(int)
tt_cores = [None] * num_dims
with tf.name_scope(name):
for i in range(num_dims):
curr_core_shape = (batch_size, tt_rank[i], shape[i],
tt_rank[i + 1])
tt_cores[i] = tf.random.normal(curr_core_shape,
mean=mean,
stddev=stddev,
dtype=dtype)
return TensorTrainBatch(tt_cores, shape, tt_rank, batch_size)
def testFullMatrix2d(self):
np.random.seed(1)
for rank in [1, 2]:
a = np.random.rand(3, 2, 3, rank).astype(self.dtype.as_numpy_dtype)
b = np.random.rand(3, rank, 4, 5).astype(self.dtype.as_numpy_dtype)
tt_cores = (a.reshape(3, 1, 2, 3, rank), b.reshape((3, rank, 4, 5, 1)))
# Basically do full by hand.
desired = np.einsum('oijb,obkl->oijkl', a, b)
desired = desired.reshape((3, 2, 3, 4, 5))
desired = desired.transpose((0, 1, 3, 2, 4))
desired = desired.reshape((3, 2 * 4, 3 * 5))
with self.test_session():
tf_mat = TensorTrainBatch(tt_cores)
actual = ops.full(tf_mat)
self.assertAllClose(desired, actual.eval())
def assign(ref, value, validate_shape=None, use_locking=None, name=None):
new_cores = []
if name is None:
name = ''
with tf.variable_scope(name):
for i in range(ref.ndims()):
new_cores.append(tf.assign(ref.tt_cores[i], value.tt_cores[i],
use_locking=use_locking))
if isinstance(value, TensorTrainBatch):
return TensorTrainBatch(new_cores, value.get_raw_shape(),
value.get_tt_ranks(), value.batch_size,
convert_to_tensors=False)
else:
return TensorTrain(new_cores, value.get_raw_shape(),
value.get_tt_ranks(),
convert_to_tensors=False)
def testFullMatrix3d(self):
np.random.seed(1)
for rank in [1, 2]:
a = np.random.rand(3, 2, 3, rank).astype(self.dtype.as_numpy_dtype)
b = np.random.rand(3, rank, 4, 5,
rank).astype(self.dtype.as_numpy_dtype)
c = np.random.rand(3, rank, 2, 2).astype(self.dtype.as_numpy_dtype)
tt_cores = (a.reshape(3, 1, 2, 3,
rank), b.reshape(3, rank, 4, 5, rank),
c.reshape(3, rank, 2, 2, 1))
# Basically do full by hand.
desired = np.einsum('oija,oaklb,obpq->oijklpq', a, b, c)
desired = desired.reshape((3, 2, 3, 4, 5, 2, 2))
desired = desired.transpose((0, 1, 3, 5, 2, 4, 6))
desired = desired.reshape((3, 2 * 4 * 2, 3 * 5 * 2))
tf_mat = TensorTrainBatch(tt_cores)
actual = self.evaluate(ops.full(tf_mat))
self.assertAllClose(desired, actual)
def testFullMatrix3d(self):
np.random.seed(1)
for rank in [1, 2]:
a = np.random.rand(3, 2, 3, rank).astype(np.float32)
b = np.random.rand(3, rank, 4, 5, rank).astype(np.float32)
c = np.random.rand(3, rank, 2, 2).astype(np.float32)
tt_cores = (a.reshape(3, 1, 2, 3,
rank), b.reshape(3, rank, 4, 5, rank),
c.reshape(3, rank, 2, 2, 1))
# Basically do full by hand.
desired = np.einsum('oija,oaklb,obpq->oijklpq', a, b, c)
desired = desired.reshape((3, 2, 3, 4, 5, 2, 2))
desired = desired.transpose((0, 1, 3, 5, 2, 4, 6))
desired = desired.reshape((3, 2 * 4 * 2, 3 * 5 * 2))
with self.test_session():
tf_mat = TensorTrainBatch(tt_cores)
actual = ops.full(tf_mat)
self.assertAllClose(desired, actual.eval())
def random_tensor_batch(shape, tt_rank=2, batch_size=1):
"""Generate a batch of random TT-tensors of given shape.
Args:
shape: array representing the shape of the future tensor
tt_rank: a number or a (d+1)-element array with ranks.
batch_size: an integer.
Returns:
TensorTrainBatch containing a TT-tensor
"""
# TODO: good distribution to init training.
# TODO: support shape and tt_ranks as TensorShape?.
# TODO: support None as a dimension.
shape = np.array(shape)
tt_rank = np.array(tt_rank)
if len(shape.shape) != 1:
raise ValueError('shape should be 1d array')
if np.any(shape < 1):
raise ValueError('all elements in `shape` should be positive')
if np.any(tt_rank < 1):
raise ValueError('`rank` should be positive')
if tt_rank.size != 1 and tt_rank.size != (shape.size + 1):
raise ValueError('`rank` array has inappropriate size')
if batch_size < 1:
raise ValueError('Batch size should be positive')
num_dims = shape.size
if tt_rank.size == 1:
tt_rank = tt_rank * np.ones(num_dims - 1)
tt_rank = np.insert(tt_rank, 0, 1)
tt_rank = np.append(tt_rank, 1)
# TODO: check that ints?
shape = shape.astype(int)
batch_size = int(batch_size)
tt_rank_ext = tt_rank.astype(int)
# TODO: variable (name?) scope.
tt_cores = [None] * num_dims
for i in range(num_dims):
curr_core_shape = (batch_size, tt_rank_ext[i], shape[i],
tt_rank_ext[i + 1])
tt_cores[i] = tf.random_normal(curr_core_shape)
return TensorTrainBatch(tt_cores, shape, tt_rank_ext, batch_size)
def expand_batch_dim(tt):
"""Creates a 1-element TensorTrainBatch from a TensorTrain.
Args:
tt: TensorTrain or TensorTrainBatch.
Returns:
TensorTrainBatch
"""
if hasattr(tt, 'batch_size'):
return tt
else:
from t3f.tensor_train_batch import TensorTrainBatch
tt_cores = []
for core_idx in range(tt.ndims()):
tt_cores.append(tf.expand_dims(tt.tt_cores[core_idx], 0))
return TensorTrainBatch(tt_cores,
tt.get_raw_shape(),
tt.get_tt_ranks(),
batch_size=1)
def testCholesky(self):
# Tests the cholesky function
np.random.seed(8)
# generating two symmetric positive-definite tt-cores
L_1 = np.tril(np.random.normal(scale=2., size=(4, 2, 2)))
L_2 = np.tril(np.random.normal(scale=2., size=(4, 3, 3)))
K_1 = np.einsum('ijk,ilk->ijl', L_1, L_1)
K_2 = np.einsum('ijk,ilk->ijl', L_2, L_2)
initializer = TensorTrainBatch(
[K_1[:, None, :, :, None], K_2[:, None, :, :, None]],
tt_ranks=7 * [1])
kron_mat_batch = variables.get_variable('kron_mat_batch',
initializer=initializer)
init_op = tf.global_variables_initializer()
with self.test_session() as sess:
sess.run(init_op)
desired = np.linalg.cholesky(ops.full(kron_mat_batch).eval())
actual = ops.full(kr.cholesky(kron_mat_batch)).eval()
self.assertAllClose(desired, actual)
def expand_batch_dim(tt, name='t3f_expand_batch_dim'):
"""Creates a 1-element TensorTrainBatch from a TensorTrain.
Args:
tt: TensorTrain or TensorTrainBatch.
name: string, name of the Op.
Returns:
TensorTrainBatch
"""
with tf.name_scope(name, values=tt.tt_cores):
if hasattr(tt, 'batch_size'):
return tt
else:
from t3f.tensor_train_batch import TensorTrainBatch
tt_cores = []
for core_idx in range(tt.ndims()):
tt_cores.append(tf.expand_dims(tt.tt_cores[core_idx], 0))
return TensorTrainBatch(tt_cores,
tt.get_raw_shape(),
tt.get_tt_ranks(),
batch_size=1)
def concat_along_batch_dim(tt_list):
"""Concat all TensorTrainBatch objects along batch dimension.
Args:
tt_list: a list of TensorTrainBatch objects.
Returns:
TensorTrainBatch
"""
ndims = tt_list[0].ndims()
if isinstance(tt_list, TensorTrainBase):
# Not a list but just one element, nothing to concat.
return tt_list
for batch_idx in range(len(tt_list)):
if not isinstance(tt_list[batch_idx], TensorTrainBatch):
raise ValueError('All objects in the list should be TTBatch objects, got '
'%s' % tt_list[batch_idx])
for batch_idx in range(1, len(tt_list)):
if tt_list[batch_idx].get_raw_shape() != tt_list[0].get_raw_shape():
raise ValueError('Shapes of all TT-batch objects should coincide, got %s '
'and %s' % (tt_list[0].get_raw_shape(),
tt_list[batch_idx].get_raw_shape()))
if tt_list[batch_idx].get_tt_ranks() != tt_list[0].get_tt_ranks():
raise ValueError('TT-ranks of all TT-batch objects should coincide, got '
'%s and %s' % (tt_list[0].get_tt_ranks(),
tt_list[batch_idx].get_tt_ranks()))
res_cores = []
for core_idx in range(ndims):
curr_core = tf.concat([tt.tt_cores[core_idx] for tt in tt_list], axis=0)
res_cores.append(curr_core)
batch_size = sum([tt.batch_size for tt in tt_list])
return TensorTrainBatch(res_cores, tt_list[0].get_raw_shape(),
tt_list[0].get_tt_ranks(), batch_size)
def transpose(tt_matrix):
"""Transpose a TT-matrix or a batch of TT-matrices.
Args:
tt_matrix: `TensorTrain` or `TensorTrainBatch` object containing a TT-matrix
(or a batch of TT-matrices).
Returns:
`TensorTrain` or `TensorTrainBatch` object containing a transposed TT-matrix
(or a batch of TT-matrices).
Raises:
ValueError if the argument is not a TT-matrix.
"""
if not isinstance(tt_matrix,
TensorTrainBase) or not tt_matrix.is_tt_matrix():
raise ValueError('The argument should be a TT-matrix.')
transposed_tt_cores = []
for core_idx in range(tt_matrix.ndims()):
curr_core = tt_matrix.tt_cores[core_idx]
if isinstance(tt_matrix, TensorTrain):
transposed_tt_cores.append(tf.transpose(curr_core, (0, 2, 1, 3)))
else:
# TensorTrainBatch.
transposed_tt_cores.append(tf.transpose(curr_core,
(0, 1, 3, 2, 4)))
tt_matrix_shape = tt_matrix.get_raw_shape()
transposed_shape = tt_matrix_shape[1], tt_matrix_shape[0]
tt_ranks = tt_matrix.get_tt_ranks()
if isinstance(tt_matrix, TensorTrain):
return TensorTrain(transposed_tt_cores, transposed_shape, tt_ranks)
else:
batch_size = tt_matrix.batch_size
return TensorTrainBatch(transposed_tt_cores, transposed_shape,
tt_ranks, batch_size)
def get_variable(name,
dtype=None,
initializer=None,
regularizer=None,
trainable=True,
collections=None,
caching_device=None,
validate_shape=True):
"""Returns TensorTrain object with tf.Variables as the TT-cores.
Args:
name: The name of the new or existing TensorTrain variable.
Used to name the TT-cores.
dtype: Type of the new or existing TensorTrain variable TT-cores (defaults
to DT_FLOAT).
initializer: TensorTrain or TensorTrainBatch, initializer for the variable
if one is created.
regularizer: A (TensorTrain -> Tensor or None) function; the result of
applying it on a newly created variable will be added to the collection
GraphKeys.REGULARIZATION_LOSSES and can be used for regularization.
trainable: If True also add the variable to the graph collection
GraphKeys.TRAINABLE_VARIABLES (see tf.Variable).
collections: List of graph collections keys to add the Variables
(underlying TT-cores). Defaults to [GraphKeys.GLOBAL_VARIABLES]
(see tf.Variable).
caching_device: Optional device string or function describing where
the Variable should be cached for reading. Defaults to the Variable's
device. If not None, caches on another device. Typical use is to cache
on the device where the Ops using the Variable reside, to deduplicate
copying through Switch and other conditional statements.
validate_shape: If False, allows the variable to be initialized with a value
of unknown shape. If True, the default, the shape of initial_value must be
known.
Returns:
The created or existing `TensorTrain` object with tf.Variables TT-cores.
Raises:
`ValueError`: when creating a new variable and shape is not declared, when
violating reuse during variable creation, or when initializer dtype and
dtype don't match. Reuse is set inside variable_scope.
"""
# TODO: support validate shape: check that the tensor dimensions are correct,
# but ignore the ranks.
# TODO: add validate ranks flag.
reuse = tf.get_variable_scope().reuse
if not reuse and initializer is None:
raise ValueError(
'Scope reuse is False and initializer is not provided.')
variable_cores = []
if reuse and not utils.in_eager_mode():
# Find an existing variable in the collection.
path = tf.get_variable_scope().name
if path != '' and path[-1] != '/':
path += '/'
path += name
found_v = None
for v in tf.get_collection('TensorTrainVariables'):
if v.name == path:
found_v = v
break
if found_v is None:
raise ValueError(
'ValueError: Variable %s does not exist, or was not '
'created with t3f.get_tt_variable(). Did you mean to '
'set reuse=None in VarScope?' % name)
with tf.variable_scope(name):
# Try to get the first core through tf.get_variable to check that we don't
# violate reuse: it will raise a ValueError otherwise.
tf.get_variable('core_0', dtype=dtype)
return found_v
else:
# Create new variable.
with tf.variable_scope(name):
num_dims = initializer.ndims()
for i in range(num_dims):
curr_core_var = tf.get_variable(
'core_%d' % i,
initializer=initializer.tt_cores[i],
dtype=dtype,
trainable=trainable,
collections=collections,
caching_device=caching_device)
variable_cores.append(curr_core_var)
if isinstance(initializer, TensorTrain):
v = TensorTrain(variable_cores,
initializer.get_raw_shape(),
initializer.get_tt_ranks(),
convert_to_tensors=False)
else:
v = TensorTrainBatch(variable_cores,
initializer.get_raw_shape(),
initializer.get_tt_ranks(),
initializer.batch_size,
convert_to_tensors=False)
# Add the create TensorTrain object into a collection so that we can
# retrieve it in the future by get_tt_variable('name').
tf.add_to_collection('TensorTrainVariables', v)
# Run the regularizer if requested and save the resulting loss.
if regularizer:
with tf.name_scope(name + "/Regularizer/"):
loss = regularizer(v)
if loss is not None:
tf.logging.vlog(
1, "Applied regularizer to %s and added the result %s "
"to REGULARIZATION_LOSSES.", v.name, loss.name)
tf.add_to_collection(tf.GraphKeys.REGULARIZATION_LOSSES, loss)
return v
def _orthogonalize_batch_tt_cores_left_to_right(tt):
"""Orthogonalize TT-cores of a batch TT-object in the left to right order.
Args:
tt: TensorTrainBatch.
Returns:
TensorTrainBatch
"""
# Left to right orthogonalization.
ndims = tt.ndims()
raw_shape = shapes.lazy_raw_shape(tt)
tt_ranks = shapes.lazy_tt_ranks(tt)
next_rank = tt_ranks[0]
batch_size = shapes.lazy_batch_size(tt)
# Copy cores references so we can change the cores.
tt_cores = list(tt.tt_cores)
for core_idx in range(ndims - 1):
curr_core = tt_cores[core_idx]
# TT-ranks could have changed on the previous iteration, so `tt_ranks` can
# be outdated for the current TT-rank, but should be valid for the next
# TT-rank.
curr_rank = next_rank
next_rank = tt_ranks[core_idx + 1]
if tt.is_tt_matrix():
curr_mode_left = raw_shape[0][core_idx]
curr_mode_right = raw_shape[1][core_idx]
curr_mode = curr_mode_left * curr_mode_right
else:
curr_mode = raw_shape[0][core_idx]
qr_shape = (batch_size, curr_rank * curr_mode, next_rank)
curr_core = tf.reshape(curr_core, qr_shape)
curr_core, triang = tf.qr(curr_core)
if triang.get_shape().is_fully_defined():
triang_shape = triang.get_shape().as_list()
else:
triang_shape = tf.shape(triang)
# The TT-rank could have changed: if qr_shape is e.g. 4 x 10, than q would
# be of size 4 x 4 and r would be 4 x 10, which means that the next rank
# should be changed to 4.
next_rank = triang_shape[1]
if tt.is_tt_matrix():
new_core_shape = (batch_size, curr_rank, curr_mode_left,
curr_mode_right, next_rank)
else:
new_core_shape = (batch_size, curr_rank, curr_mode, next_rank)
tt_cores[core_idx] = tf.reshape(curr_core, new_core_shape)
next_core = tf.reshape(tt_cores[core_idx + 1],
(batch_size, triang_shape[2], -1))
tt_cores[core_idx + 1] = tf.matmul(triang, next_core)
if tt.is_tt_matrix():
last_core_shape = (batch_size, next_rank, raw_shape[0][-1],
raw_shape[1][-1], 1)
else:
last_core_shape = (batch_size, next_rank, raw_shape[0][-1], 1)
tt_cores[-1] = tf.reshape(tt_cores[-1], last_core_shape)
# TODO: infer the tt_ranks.
return TensorTrainBatch(tt_cores,
tt.get_raw_shape(),
batch_size=batch_size)
