Dear ITensor Teams
Thank you for developing such user-friendly and efficient code.
I am currently working on optimizing matrix product states (MPS)/ tensor trains (TT) using Zygote for automatic differentiation. Specifically, I am trying to optimize each element of the MPS/TT tensors. However, I have encountered an issue when attempting to compress the bond dimensions within the loss function.
In the provided code, I am using ITensors.truncate inside the ttfit function, which serves as the loss function. Unfortunately, this results in the error: “no method matching iterate(::Nothing)”. I suspect this error arises because ITensors.truncate might return nothing under certain conditions.
My main question is: What would be a suitable approach to perform bond dimension compression within the loss function, given the constraints of automatic differentiation with Zygote?
I would be grateful for any insights, suggestions, or best practices you could share regarding this problem.
I have attached a minimum code to reproduce my issue below.
using ITensors
using Random
using OptimKit
using Zygote
ITensors.disable_warn_order()
R = 8 #
D = 2 #
fit_itr = 10 #
sitedim = 2
sites = siteinds(sitedim, R)
#Make random MPS
function make_random_mps(sitedims::Vector{Int64}, linkdims::Vector{Int64}, nsites::Int64)::Vector{Array{Float64, 3}}
res = []
for i in 1:nsites
if i == 1
push!(res, randn(1, sitedims[i], linkdims[i]))
elseif i == nsites
push!(res, randn(linkdims[end], sitedims[i], 1))
else
push!(res, randn(linkdims[i-1], sitedims[i], linkdims[i]))
end
end
return res
end
# Get bond dimension
function linkdims(tt::Vector{Array{Float64, 3}})
return [size(T, 1) for T in tt[2:end]]
end
# Convert MPS to ITensors.MPS
function ITensors.MPS(tt::Vector{Array{Float64, 3}}, sites=nothing)::MPS
N = length(tt)
localdims = [size(t, 2) for t in tt]
if sites === nothing
sites = [Index(localdims[n], "n=$n") for n in 1:N]
else
all(localdims .== [ITensors.dim(sites[i]) for i in 1:N]) ||
error("ranks are not consistent with dimension of sites")
end
linkdims_ = [1, linkdims(tt)..., 1]
links = [Index(linkdims_[l + 1], "link, l=$l") for l in 0:N]
tensors_ = [ITensor(deepcopy(tt[n]), links[n], sites[n], links[n+1]) for n in 1:N]
return MPS(tensors_)
end
# TensorTrainFit struct
struct TensorTrainFit{ValueType}
tt::Vector{Array{ValueType, 3}}
offsets::Vector{Int}
end
# Make offset, which is used to convert 1D array to MPS
function TensorTrainFit{ValueType}(tt) where {ValueType}
offsets = [0]
for n in 1:length(tt)
push!(offsets, offsets[end] + length(tt[n]))
end
return TensorTrainFit{ValueType}(tt, offsets)
end
function (obj::TensorTrainFit{ValueType})(x::Vector{ValueType}) where {ValueType}
tensors = to_tensors(obj, x) # Convert 1D array to MPS
tensors2 = ITensors.MPS(tensors, sites) # Convert to ITensors.MPS
tensors3 = ITensors.truncate(tensors2; cutoff=1e-10) # Truncate
return inner(tensors3, tensors3) # Output inner product, which is the loss function
#return inner(tensors2, tensors2) # this is working if we comment out above 2 lines
end
# Convert each core tensors of MPS to 1D array
function to_tensors(obj::TensorTrainFit{ValueType}, x::Vector{ValueType}) where {ValueType}
return [
reshape(
x[obj.offsets[n]+1:obj.offsets[n+1]],
size(obj.tt[n])
)
for n in 1:length(obj.tt)
]
end
#Convert to 1D array
flatten_(obj::Vector{Array{ValueType, 3}}) where {ValueType} = vcat(vec.(obj)...)
tt0 = make_random_mps([2 for i in 1:R], [D for i in 1:R-1], R)
ttfit = TensorTrainFit{Float64}(tt0) # Create TensorTrainFit object
x0::Vector{Float64} = flatten_(tt0) #Convert to 1D array
function loss(x)
ttfit(x)
end
optimizer = LBFGS(; maxiter=fit_itr, verbosity=1)
function loss_and_grad(x)
y, (∇,) = withgradient(loss, x)
return y, ∇
end
withgradient(loss, x0) # falied
xopt, fs, gs, niter, normgradhistory = optimize(loss_and_grad, x0, optimizer)
The version of the package I use is the following.
ITensors v0.3.68
OptimKit v0.3.1
Zygote v0.6.70
Thank you for your time.
Best,
Rihito