Automatic differentiation

Being based on pytorch, torchtt can handle automatic differentiation with respect to the TT cores.

Imports

import torch as tn
try:
    import torchtt as tntt
except:
    print('Installing torchTT...')
    %pip install git+https://github.com/ion-g-ion/torchTT
    import torchtt as tntt

First, a function to differentiate is created and some tensors:

N = [2,3,4,5]
A = tntt.randn([(n,n) for n in N],[1]+[2]*(len(N)-1)+[1])
y = tntt.randn(N,A.R)
x = tntt.ones(N)

def f(x,A,y):
    z = tntt.dot(A @ (x-y),(x-y))
    return z.norm()

In order to compute the derivative of a scalar with respect to all cores of a TT object, the AD graph recording has to be started:

tntt.grad.watch(x)

Using the torchtt.grad.grad() method, the gradient is computed:

val = f(x,A,y)
grad_cores = tntt.grad.grad(val, x)

The variable grad_cores is a list of tensors representing the derivatives of f() with resect to the individual core entries. For checking, we compute the derivative of teh function with respect to one element of the core

h = 1e-7
x1 = x.clone()
x1.cores[1][0,0,0] += h
x2 = x.clone()
x2.cores[1][0,0,0] -= h
derivative = (f(x1,A,y)-f(x2,A,y))/(2*h)
print(tn.abs(derivative-grad_cores[1][0,0,0])/tn.abs(derivative))

tensor(1.2106e-09, dtype=torch.float64, grad_fn=<DivBackward0>)

The functions torchtt.grad.grad() and torchtt.grad.watch() can take an additional list of modes core_indices as argument which decides which cores are watched and differentiaated with respect to.