"""
Manifold gradient module.
"""
import torch as tn
from torchtt._decomposition import mat_to_tt, to_tt, lr_orthogonal, round_tt, rl_orthogonal
from . import TT
from torchtt.errors import *
def _delta2cores(tt_cores, R, Sds, is_ttm = False, ortho = None):
"""
Convert the detla notation to TT.
Implements Algorithm 5.1 from "AUTOMATIC DIFFERENTIATION FOR RIEMANNIAN OPTIMIZATION ON LOW-RANK MATRIX AND TENSOR-TRAIN MANIFOLDS".
Args:
tt_cores (list[torch.tensor]): the TT cores.
R (list[int]): the rank of the tensor.
Sds (list[torch.tensor]): deltas.
is_ttm (bool, optional): is TT amtrix or not. Defaults to False.
ortho (list[list[torch.tensor]], optional): the left and right orthogonal cores of tt_cores. Defaults to None.
Returns:
list[torch.tensor]: the resulting TT cores.
"""
if ortho == None:
l_cores,_ = lr_orthogonal(tt_cores, R, is_ttm)
r_cores,_ = rl_orthogonal(tt_cores, R, is_ttm)
else:
l_cores = ortho[0]
r_cores = ortho[1]
# first
cores_new = [tn.cat((Sds[0],l_cores[0]),2 if not is_ttm else 3)]
# 2...d-1
for k in range(1,len(tt_cores)-1):
up = tn.cat((r_cores[k],tn.zeros((r_cores[k].shape),dtype = l_cores[0].dtype, device = l_cores[0].device)),2 if not is_ttm else 3)
down = tn.cat((Sds[k],l_cores[k]),2 if not is_ttm else 3)
cores_new.append(tn.cat((up,down),0))
# last
cores_new.append(tn.cat((r_cores[-1],Sds[-1]),0))
return cores_new
[docs]
def riemannian_gradient(x,func):
"""
Compute the Riemannian gradient using AD.
Args:
x (torchtt.TT): the point on the manifold where the gradient is computed.
func ([type]): function that has to be differentiated. The function takes as only argument `torchtt.TT` instances.
Returns:
torchtt.TT: the gradient projected on the tangent space of x.
"""
l_cores,_ = lr_orthogonal(x.cores, x.R, x.is_ttm)
r_cores,_ = rl_orthogonal(l_cores, x.R, x.is_ttm)
is_ttm = x.is_ttm
R = x.R
d = len(x.N)
Rs = [ r_cores[0] ]
Rs += [ x.cores[i]*0 for i in range(1,d)]
# AD part
for i in range(d):
Rs[i].requires_grad_(True)
Ghats = _delta2cores(x.cores, R, Rs, is_ttm = is_ttm,ortho = [l_cores,r_cores])
fval = func(TT(Ghats))
fval.backward()
# Sds = tape.gradient(fval, Rs)
Sds = [r.grad for r in Rs]
# print('Sds ',Sds)
# compute Sdeltas
for k in range(d-1):
D = tn.reshape(Sds[k],[-1,R[k+1]])
UL = tn.reshape(l_cores[k],[-1,R[k+1]])
D = D - UL @ (UL.T @ D)
Sds[k] = tn.reshape(D,l_cores[k].shape)
# print([tf.einsum('ijk,ijl->kl',l_cores[i],Sds[i]).numpy() for i in range(d-1)])
# delta to TT
grad_cores = _delta2cores(x.cores, R, Sds, is_ttm,ortho = [l_cores,r_cores])
return TT(grad_cores)
[docs]
def riemannian_projection(Xspace,z):
"""
Project the tensor z onto the tangent space defined at xspace
Args:
Xspace (torchtt.TT): the target where the tensor should be projected.
z (torchtt.TT): the tensor that should be projected.
Raises:
IncompatibleTypes: Both must be of same type.
Returns:
torchtt.TT: the projection.
"""
if Xspace.is_ttm != z.is_ttm:
raise IncompatibleTypes('Both must be of same type.')
is_ttm = Xspace.is_ttm
l_cores,R = lr_orthogonal(Xspace.cores, Xspace.R, Xspace.is_ttm)
r_cores,_ = rl_orthogonal(l_cores, R, Xspace.is_ttm)
d = len(Xspace.N)
N = Xspace.N
# Pleft = [tf.ones((1,1,1),dtype=Xspace.cores[0].dtype)]
Pleft = []
tmp = tn.ones((1,1),dtype=Xspace.cores[0].dtype, device = Xspace.cores[0].device)
for k in range(d-1):
if is_ttm:
tmp = tn.einsum('rs,rijR,sijS->RS',tmp,l_cores[k],z.cores[k]) # size rk x sk
else:
tmp = tn.einsum('rs,riR,siS->RS',tmp,l_cores[k],z.cores[k]) # size rk x sk
Pleft.append(tmp)
Pright = []
tmp = tn.ones((1,1), dtype = Xspace.cores[0].dtype, device = Xspace.cores[0].device)
for k in range(d-1,0,-1):
if is_ttm:
tmp = tn.einsum('RS,rijR,sijS->rs',tmp,r_cores[k],z.cores[k]) # size rk x sk
else:
tmp = tn.einsum('RS,riR,siS->rs',tmp,r_cores[k],z.cores[k]) # size rk x sk
Pright.append(tmp)
Pright = Pright[::-1]
# compute elements of the tangent space
Sds = []
for k in range(d):
if k==0:
L = tn.ones((1,1),dtype=Xspace.cores[0].dtype, device = Xspace.cores[0].device)
else:
L = Pleft[k-1]
if k==d-1:
if is_ttm:
Sds.append(tn.einsum('rs,sjiS->rjiS',L,z.cores[k]))
else:
Sds.append(tn.einsum('rs,siS->riS',L,z.cores[k]))
else:
R = Pright[k]
if is_ttm:
tmp1 = tn.einsum('rs,sijS->rijS',L,z.cores[k])
tmp2 = tn.einsum('rijR,RS->rijS',l_cores[k],tn.einsum('rs,rijR,sijS->RS',L,l_cores[k],z.cores[k]))
Sds.append(tn.einsum('rijS,RS->rijR',tmp1-tmp2,R))
else:
tmp1 = tn.einsum('rs,siS->riS',L,z.cores[k])
tmp2 = tn.einsum('riR,RS->riS',l_cores[k],tn.einsum('rs,riR,siS->RS',L,l_cores[k],z.cores[k]))
Sds.append(tn.einsum('riS,RS->riR',tmp1-tmp2,R))
# convert Sds to TT
grad_cores = _delta2cores(Xspace.cores, R, Sds, Xspace.is_ttm,ortho = [l_cores,r_cores])
return TT(grad_cores)