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)