Unexpectedly slow partial tensor train evaluation

[MFrankenbach]

Description

I have the following functions to saturate given legs of a tensor train with given values and contract the remainder to a dense tensor:

ITensor:

"""
Saturate legs of `mps` at sites `idx_pos` with indices `sidx[idx_pos]` taking values `idx` and contract result to dense ITensor.
`sidx` are ALL site indices of MPS.
"""
function eval_mps_partial(mps::MPS, sidx::Vector{<:Index}, idx::Vector{Int}, idx_pos::AbstractVector{Int})::ITensor
    res = ITensor(one(eltype(mps[1])))
    for i in eachindex(mps.data)
        ii = findfirst(isequal(i), idx_pos)
        if isnothing(ii)
            res *= mps[i]
        else
            res *= mps[i] * onehot(sidx[i] => idx[ii])
        end
    end
    return res
end

Tensor4all

function eval_mps_partial(
    tt::Tensor4all.TensorNetworks.TensorTrain,
    sidx::Vector{<:Tensor4all.Index},
    idx::Vector{Int},
    idx_pos::AbstractVector{Int},
)::Tensor4all.Tensor
    length(idx) == length(idx_pos) || error("Number of fixed values must match number of fixed positions")
    isempty(idx_pos) && return Tensor4all.TensorNetworks.to_dense(tt)

    replacements = map(Pair, sidx[idx_pos], idx)
    return Tensor4all.TensorNetworks.to_dense(
        Tensor4all.TensorNetworks.fixinds(tt, replacements...)
    )
end

Now I observed that the Tensor4all version is much slower. Even slower is the following, even though it matches the way I do it with ITensor:

function eval_mps_partial_old(
    tt::Tensor4all.TensorNetworks.TensorTrain,
    sidx::Vector{<:Tensor4all.Index},
    idx::Vector{Int},
    idx_pos::AbstractVector{Int},
)::Tensor4all.Tensor
    N = length(tt)
    fixed_mask = falses(N)
    fixed_values = zeros(Int, N)
    for (pos, value) in zip(idx_pos, idx)
        fixed_mask[pos] = true
        fixed_values[pos] = value
    end

    result = nothing
    for i in 1:N
        local_tensor = tt[i]
        if fixed_mask[i]
            basis = zeros(eltype(local_tensor), Tensor4all.dim(sidx[i]))
            basis[fixed_values[i]] = one(eltype(local_tensor))
            local_tensor = Tensor4all.contract(local_tensor, Tensor4all.Tensor(basis, [sidx[i]]))
        end
        result = isnothing(result) ? local_tensor : Tensor4all.contract(result, local_tensor)
    end
    return result::Tensor4all.Tensor
end

Steps to reproduce

Create Tensor4all/ITensor tensor trains and evaluate them with the above functions.

Actual behavior

Evaluating two equivalent tensor trains in ITensor and Tensor4all yields vastly different runtimes:

julia> @btime BubbleTeaCI.eval_mps_partial(tf_ITensor.data, tf_ITensor.site_indices, [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], tf_ITensor.continuous_legs)
  68.541 μs (796 allocations: 355.88 KiB)
ITensor ord=2 (dim=2|id=901|"idx1") (dim=3|id=328|"idx2")
NDTensors.Dense{Float64, Vector{Float64}}

julia> @btime BubbleTeaCI.eval_mps_partial(tf_Tensor4all.data, tf_Tensor4all.site_indices, [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], tf_Tensor4all.continuous_legs)
  2.084 ms (12213 allocations: 1.57 MiB)
Tensor4all.Tensor{Float64, 2}(Tensor4all.TensorHandle(Ptr{Nothing} @0x0000000106debb50), Tensor4all.Index[Index(2|idx1; plev=0), Index(3|idx2; plev=0)])

Also, Tensor4all allocates a lot more memory. Tensors are copied unnecessarily due to interfacing with Rust?

Expected behavior

Tensor4all should not be that much slower? Or am I missing something?

Tensor4all.jl version

No response

Julia version

1.11.7

OS

mac OS