def _orthogonalize_tt_cores_right_to_left(tt):
"""Orthogonalize TT-cores of a TT-object in the right to left order.
Args:
tt: TenosorTrain or a TensorTrainBatch.
Returns:
The same type as the input `tt` (TenosorTrain or a TensorTrainBatch).
"""
# Left to right orthogonalization.
ndims = tt.ndims()
raw_shape = shapes.lazy_raw_shape(tt)
tt_ranks = shapes.lazy_tt_ranks(tt)
prev_rank = tt_ranks[ndims]
# Copy cores references so we can change the cores.
tt_cores = list(tt.tt_cores)
for core_idx in range(ndims - 1, 0, -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 = prev_rank
prev_rank = tt_ranks[core_idx]
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 = (prev_rank, curr_mode * curr_rank)
curr_core = tf.reshape(curr_core, qr_shape)
curr_core, triang = tf.qr(tf.transpose(curr_core))
curr_core = tf.transpose(curr_core)
triang = tf.transpose(triang)
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.
prev_rank = triang_shape[1]
if tt.is_tt_matrix():
new_core_shape = (prev_rank, curr_mode_left, curr_mode_right,
curr_rank)
else:
new_core_shape = (prev_rank, curr_mode, curr_rank)
tt_cores[core_idx] = tf.reshape(curr_core, new_core_shape)
prev_core = tf.reshape(tt_cores[core_idx - 1], (-1, triang_shape[0]))
tt_cores[core_idx - 1] = tf.matmul(prev_core, triang)
if tt.is_tt_matrix():
first_core_shape = (1, raw_shape[0][0], raw_shape[1][0], prev_rank)
else:
first_core_shape = (1, raw_shape[0][0], prev_rank)
tt_cores[0] = tf.reshape(tt_cores[0], first_core_shape)
# TODO: infer the tt_ranks.
return TensorTrain(tt_cores, tt.get_raw_shape())
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 eye(shape, dtype=tf.float32, name='t3f_eye'):
"""Creates an identity TT-matrix.
Args:
shape: array which defines the shape of the matrix row and column
indices.
dtype: [tf.float32] dtype of the resulting matrix.
name: string, name of the Op.
Returns:
TensorTrain containing an identity TT-matrix of size
np.prod(shape) x np.prod(shape)
"""
shape = np.array(shape)
# In this special case shape is in the same format as in the TT-tensor case
_validate_input_parameters(is_tensor=True, shape=shape)
num_dims = shape.size
tt_ranks = np.ones(num_dims + 1, dtype=np.int)
with tf.name_scope(name):
tt_cores = num_dims * [None]
for i in range(num_dims):
curr_core_shape = (1, shape[i], shape[i], 1)
tt_cores[i] = tf.reshape(tf.eye(shape[i], dtype=dtype),
curr_core_shape)
true_shape = np.vstack([shape, shape])
return TensorTrain(tt_cores, true_shape, tt_ranks)
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 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 _batch_dim_getitem(self, element_spec):
"""__getitem__ when provided only one (batch) index.
Examples:
a[1]
a[1:3]
"""
# This object index is specified exactly and we want to collapse the
# batch_size axis, i.e. return a TensorTrain instead of a TensorTrainBatch.
do_collapse_batch_dim = self._do_collapse_dim(element_spec)
new_tt_cores = []
for core_idx in range(self.ndims()):
curr_core = self.tt_cores[core_idx]
if self.is_tt_matrix():
new_tt_cores.append(curr_core[element_spec, :, :, :, :])
else:
new_tt_cores.append(curr_core[element_spec, :, :, :])
if do_collapse_batch_dim:
# This index is specified exactly and we want to collapse the batch_size
# axis, i.e. return a TensorTrain instead of a TensorTrainBatch.
return TensorTrain(new_tt_cores, self.get_raw_shape(),
self.get_tt_ranks())
else:
batch_size = new_tt_cores[0].get_shape()[0].value
return TensorTrainBatch(new_tt_cores, self.get_raw_shape(),
self.get_tt_ranks(), batch_size)
def matrix_with_random_cores(shape,
tt_rank=2,
mean=0.,
stddev=1.,
dtype=tf.float32,
name='t3f_matrix_with_random_cores'):
"""Generate a TT-matrix 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_with_random_cores([[2, 2, 2], None])
and
matrix_with_random_cores([None, [2, 2, 2]])
will create an 8-element column and row vectors correspondingly.
tt_rank: a number or a (d+1)-element array with ranks.
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 matrix.
name: string, name of the Op.
Returns:
TensorTrain containing a TT-matrix 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)
num_dims = shape[0].size
if tt_rank.size == 1:
tt_rank = tt_rank * np.ones(num_dims - 1, dtype=np.int)
tt_rank = np.concatenate([[1], 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 = (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,
dtype=dtype)
return TensorTrain(tt_cores, shape, tt_rank)
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 testHalfKnownRanksTTMatmul(self):
# Tests tt_tt_matmul for the case when one matrice has known ranks
# and the other one doesn't
np.random.seed(1)
K_1 = tf.placeholder(self.dtype, (1, 2, 2, None))
K_2 = tf.placeholder(self.dtype, (None, 3, 3, 1))
tt_mat_known_ranks = TensorTrain([K_1, K_2], tt_ranks=[1, 3, 1])
tt_mat = TensorTrain([K_1, K_2])
res_actual = ops.full(ops.matmul(tt_mat_known_ranks, tt_mat))
res_desired = tf.matmul(ops.full(tt_mat_known_ranks), ops.full(tt_mat))
np.random.seed(1)
K_1_val = np.random.rand(1, 2, 2, 3)
K_2_val = np.random.rand(3, 3, 3, 1)
with self.test_session() as sess:
res_actual_val = sess.run(res_actual, {K_1: K_1_val, K_2: K_2_val})
res_desired_val = sess.run(res_desired, {K_1: K_1_val, K_2: K_2_val})
self.assertAllClose(res_desired_val, res_actual_val)
def _full_getitem(self, slice_spec):
"""__getitem__ when provided full index of length ndims + 1.
Examples:
a = t3f.random_tensor_batch((2, 3, 4), batch_size=5)
a[:3, 1:2, 4, :]
"""
if len(slice_spec) != self.ndims() + 1:
raise ValueError('Expected %d indices, got %d' % (self.ndims() + 1,
len(slice_spec)))
# This object index is specified exactly and we want to collapse the
# batch_size axis, i.e. return a TensorTrain instead of a TensorTrainBatch.
do_collapse_batch_dim = self._do_collapse_dim(slice_spec[0])
remainder = None
new_tt_cores = []
for core_idx in range(self.ndims()):
curr_core = self.tt_cores[core_idx]
if self.is_tt_matrix():
raise NotImplementedError
else:
sliced_core = curr_core[slice_spec[0], :, slice_spec[core_idx + 1], :]
do_collapse_curr_dim = self._do_collapse_dim(slice_spec[core_idx + 1])
if do_collapse_curr_dim:
# This index is specified exactly and we want to collapse this axis.
if remainder is None:
remainder = sliced_core
else:
if do_collapse_batch_dim:
remainder = tf.einsum('ab,bd->ad', remainder, sliced_core)
else:
remainder = tf.einsum('oab,obd->oad', remainder, sliced_core)
else:
if remainder is not None:
# Add reminder from the previous collapsed cores to the current
# core.
if do_collapse_batch_dim:
sliced_core = tf.einsum('ab,bid->aid', remainder, sliced_core)
else:
sliced_core = tf.einsum('oab,obid->oaid', remainder,
sliced_core)
remainder = None
new_tt_cores.append(sliced_core)
if remainder is not None:
# The reminder obtained from collapsing the last cores.
if do_collapse_batch_dim:
new_tt_cores[-1] = tf.einsum('aib,bd->aid', new_tt_cores[-1],
remainder)
else:
new_tt_cores[-1] = tf.einsum('oaib,obd->oaid', new_tt_cores[-1],
remainder)
remainder = None
# TODO: infer the output ranks and shape.
if do_collapse_batch_dim:
return TensorTrain(new_tt_cores)
else:
return TensorTrainBatch(new_tt_cores)
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(10, rank).astype(self.dtype.as_numpy_dtype)
b = np.random.rand(rank, 9).astype(self.dtype.as_numpy_dtype)
tt_cores = (a.reshape(1, 10, rank), b.reshape(rank, 9, 1))
desired = np.dot(a, b)
tf_tens = TensorTrain(tt_cores)
actual = self.evaluate(ops.full(tf_tens))
self.assertAllClose(desired, actual)
def random_matrix(shape, tt_rank=2):
"""Generate a random TT-matrix 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([[2, 2, 2], None])
and
random_matrix([None, [2, 2, 2]])
will create an 8-element column and row vectors correspondingly.
tt_rank: a number or a (d+1)-element array with ranks.
Returns:
TensorTrain containing a TT-matrix 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')
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 = (tt_rank[i], shape[0][i], shape[1][i],
tt_rank[i + 1])
tt_cores[i] = tf.random_normal(curr_core_shape)
return TensorTrain(tt_cores, shape, tt_rank)
def testProject(self):
# Compare our projection with the results obtained (and precomputed) from
# tt.riemannian.project which is well tested.
tangent_tens_cores = ([[[-0.42095269, 0.02130842],
[-0.4181081 , 0.42945687],
[ 0.45972439, -0.4525616 ],
[-0.17159869, -0.14505528]]], [[[ 0.23344421],
[ 0.81480049],
[-0.92385135]],
[[-0.19279465],
[ 0.524976 ],
[-0.40149197]]])
convert = lambda t: np.array(t, dtype=self.dtype.as_numpy_dtype)
tangent_tens_cores = list([convert(t) for t in tangent_tens_cores])
tangent_tens = TensorTrain(tangent_tens_cores, (4, 3), (1, 2, 1))
tens_cores = ([[[-1.01761142, 0.36075896, -0.2493624 ],
[-0.99896565, -1.12685474, 1.02832458],
[ 1.08739724, -0.6537435 , 1.99975537],
[ 0.35128005, 0.40395104, -0.16790072]]], [[[ 0.34105142],
[ 0.10371947],
[-1.76464904]],
[[ 0.32639758],
[-1.27843174],
[-1.41590327]],
[[ 0.76616274],
[ 0.6577514 ],
[ 2.13703185]]])
tens_cores = list([convert(t) for t in tens_cores])
tens = TensorTrain(tens_cores, (4, 3), (1, 3, 1))
desired_projection = [[-0.67638254, -1.17163914, 0.29850939],
[-1.66479093, -0.99003251, 2.46629195],
[-0.04847773, -0.72908174, 0.20142675],
[ 0.34431125, -0.20935516, -1.15864246]]
proj = riemannian.project_sum(tens, tangent_tens)
proj_full = ops.full(proj)
with self.test_session() as sess:
proj_v = proj_full.eval()
self.assertAllClose(desired_projection, proj_v)
self.assertEqual(self.dtype.as_numpy_dtype, proj_v.dtype)
def testFullTensor2d(self):
np.random.seed(1)
for rank in [1, 2]:
a = np.random.rand(10, rank).astype(np.float32)
b = np.random.rand(rank, 9).astype(np.float32)
tt_cores = (a.reshape(1, 10, rank), b.reshape(rank, 9, 1))
desired = np.dot(a, b)
with self.test_session():
tf_tens = TensorTrain(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 testCastIntFloat(self):
# Tests cast function from int to float for matrices.
np.random.seed(1)
K_1 = np.random.randint(0, high=100, size=(1, 2, 2, 2))
K_2 = np.random.randint(0, high=100, size=(2, 3, 3, 2))
K_3 = np.random.randint(0, high=100, size=(2, 2, 2, 1))
tt_int = TensorTrain([K_1, K_2, K_3], tt_ranks=[1, 2, 2, 1])
casted = ops.cast(tt_int, self.dtype)
casted_val = self.evaluate(ops.full(casted))
self.assertEqual(self.dtype, casted.dtype)
self.assertTrue(self.dtype, casted_val.dtype)
def testFullTensor3d(self):
np.random.seed(1)
for rank_1 in [1, 2]:
a = np.random.rand(10, rank_1).astype(self.dtype.as_numpy_dtype)
b = np.random.rand(rank_1, 9, 3).astype(self.dtype.as_numpy_dtype)
c = np.random.rand(3, 8).astype(self.dtype.as_numpy_dtype)
tt_cores = (a.reshape(1, 10, rank_1), b, c.reshape((3, 8, 1)))
# Basically do full by hand.
desired = a.dot(b.reshape((rank_1, -1)))
desired = desired.reshape((-1, 3)).dot(c)
desired = desired.reshape(10, 9, 8)
tf_tens = TensorTrain(tt_cores)
actual = self.evaluate(ops.full(tf_tens))
self.assertAllClose(desired, actual)
def testCastIntFloat(self):
# Tests cast function from int to float for matrices.
np.random.seed(1)
K_1 = np.random.randint(0, high=100, size=(1, 2, 2, 2))
K_2 = np.random.randint(0, high=100, size=(2, 3, 3, 2))
K_3 = np.random.randint(0, high=100, size=(2, 2, 2, 1))
tt_int = TensorTrain([K_1, K_2, K_3], tt_ranks=[1, 2, 2, 1])
tt_int_batch = shapes.expand_batch_dim(tt_int)
with self.test_session() as sess:
casted = ops.cast(tt_int_batch, self.dtype)
casted_val = sess.run(ops.full(casted))
self.assertEqual(self.dtype, casted.dtype)
self.assertTrue(self.dtype, casted_val.dtype)
def testFullMatrix2d(self):
np.random.seed(1)
for rank in [1, 2]:
a = np.random.rand(2, 3, rank).astype(self.dtype.as_numpy_dtype)
b = np.random.rand(rank, 4, 5).astype(self.dtype.as_numpy_dtype)
tt_cores = (a.reshape(1, 2, 3, rank), b.reshape((rank, 4, 5, 1)))
# Basically do full by hand.
desired = a.reshape((-1, rank)).dot(b.reshape((rank, -1)))
desired = desired.reshape((2, 3, 4, 5))
desired = desired.transpose((0, 2, 1, 3))
desired = desired.reshape((2 * 4, 3 * 5))
tf_mat = TensorTrain(tt_cores)
actual = self.evaluate(ops.full(tf_mat))
self.assertAllClose(desired, actual)
def testCastIntFloat(self):
# Tests cast function from int to float for matrices.
np.random.seed(1)
K_1 = np.random.randint(0, high=100, size=(1, 2, 2, 2))
K_2 = np.random.randint(0, high=100, size=(2, 3, 3, 2))
K_3 = np.random.randint(0, high=100, size=(2, 2, 2, 1))
tt_int = TensorTrain([K_1, K_2, K_3], tt_ranks=[1, 2, 2, 1])
with self.test_session() as sess:
for dtype in [tf.float16, tf.float32, tf.float64]:
casted = ops.cast(tt_int, dtype)
casted_val = sess.run(ops.full(casted))
self.assertEqual(dtype, casted.dtype)
self.assertTrue(dtype, casted_val.dtype)
def testUnknownRanksTTMatmul(self):
# Tests tt_tt_matmul for matrices with unknown ranks
K_1 = tf.placeholder(self.dtype, (1, 2, 2, None))
K_2 = tf.placeholder(self.dtype, (None, 3, 3, 1))
tt_mat = TensorTrain([K_1, K_2])
res_actual = ops.full(ops.matmul(tt_mat, tt_mat))
res_desired = tf.matmul(ops.full(tt_mat), ops.full(tt_mat))
np.random.seed(1)
K_1_val = np.random.rand(1, 2, 2, 2)
K_2_val = np.random.rand(2, 3, 3, 1)
with self.test_session() as sess:
res_actual_val = sess.run(res_actual, {K_1: K_1_val, K_2: K_2_val})
res_desired_val = sess.run(res_desired, {K_1: K_1_val, K_2: K_2_val})
self.assertAllClose(res_desired_val, res_actual_val)
def testFullTensor3d(self):
np.random.seed(1)
for rank_1 in [1, 2]:
a = np.random.rand(10, rank_1).astype(np.float32)
b = np.random.rand(rank_1, 9, 3).astype(np.float32)
c = np.random.rand(3, 8).astype(np.float32)
tt_cores = (a.reshape(1, 10, rank_1), b, c.reshape((3, 8, 1)))
# Basically do full by hand.
desired = a.dot(b.reshape((rank_1, -1)))
desired = desired.reshape((-1, 3)).dot(c)
desired = desired.reshape(10, 9, 8)
with self.test_session():
tf_tens = TensorTrain(tt_cores)
actual = ops.full(tf_tens)
self.assertAllClose(desired, actual.eval())
def testFullMatrix2d(self):
np.random.seed(1)
for rank in [1, 2]:
a = np.random.rand(2, 3, rank).astype(np.float32)
b = np.random.rand(rank, 4, 5).astype(np.float32)
tt_cores = (a.reshape(1, 2, 3, rank), b.reshape((rank, 4, 5, 1)))
# Basically do full by hand.
desired = a.reshape((-1, rank)).dot(b.reshape((rank, -1)))
desired = desired.reshape((2, 3, 4, 5))
desired = desired.transpose((0, 2, 1, 3))
desired = desired.reshape((2 * 4, 3 * 5))
with self.test_session():
tf_mat = TensorTrain(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 tensor_with_random_cores(shape,
tt_rank=2,
mean=0.,
stddev=1.,
dtype=tf.float32,
name='t3f_tensor_with_random_cores'):
"""Generate a TT-tensor of the 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 the desired ranks.
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:
TensorTrain 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)
_validate_input_parameters(is_tensor=True, shape=shape, tt_rank=tt_rank)
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 = (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 TensorTrain(tt_cores, shape, tt_rank)
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=(2, 2)))
L_2 = np.tril(np.random.normal(scale=2., size=(3, 3)))
K_1 = L_1.dot(L_1.T)
K_2 = L_2.dot(L_2.T)
K = np.kron(K_1, K_2)
initializer = TensorTrain(
[K_1[None, :, :, None], K_2[None, :, :, None]], tt_ranks=7 * [1])
kron_mat = variables.get_variable('kron_mat', initializer=initializer)
init_op = tf.global_variables_initializer()
with self.test_session() as sess:
sess.run(init_op)
desired = np.linalg.cholesky(K)
actual = ops.full(kr.cholesky(kron_mat)).eval()
self.assertAllClose(desired, actual)
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=(2, 2)))
L_1 = L_1.astype(self.dtype.as_numpy_dtype)
L_2 = np.tril(np.random.normal(scale=2., size=(3, 3)))
L_2 = L_2.astype(self.dtype.as_numpy_dtype)
K_1 = L_1.dot(L_1.T)
K_2 = L_2.dot(L_2.T)
K = np.kron(K_1, K_2)
initializer = TensorTrain(
[K_1[None, :, :, None], K_2[None, :, :, None]], tt_ranks=7 * [1])
kron_mat = variables.get_variable('kron_mat', initializer=initializer)
init_op = tf.compat.v1.global_variables_initializer()
self.evaluate(init_op)
desired = np.linalg.cholesky(K)
actual = self.evaluate(ops.full(kr.cholesky(kron_mat)))
self.assertAllClose(desired, actual, atol=1e-5, rtol=1e-5)
def tensor_zeros(shape):
"""Generate TT-tensor of the given shape with all entries equal to 0.
Args:
shape: array representing the shape of the future tensor
Returns:
TensorTrain object containing a TT-tensor
"""
shape = np.array(shape)
_validate_input_parameters(is_tensor=True, shape=shape)
num_dims = shape.size
tt_rank = np.ones(num_dims + 1)
tt_cores = num_dims * [None]
for i in range(num_dims):
curr_core_shape = (1, shape[i], 1)
tt_cores[i] = tf.zeros(curr_core_shape)
return TensorTrain(tt_cores, shape, tt_rank)