Allocation overhead on multiple replications of Quantum TNs

[mofeing]

I had to implement Enzyme.onehot to make forward-differentiation work with Enzyme/Reactant.

Thanks to the replace, we can reuse the arrays that are unchanged when doing multiple replications + minor modifications of a TensorNetwork like in the following code:

function Enzyme.onehot(x::MPS)
    x0 = zero(x)
    return collect(Iterators.flatten(map(tensors(x0)) do t0
        return map(Enzyme.onehot(t0)) do t1
            replace(x0, t0 => t1)
        end
    end))
end

The number of allocations should be in the order of the number of elements, which for my case is currently a 10 site MPS of physical dim = 2 and max bond dim = 32... so 2778 elements. However, @time reveals that the allocation count is in the order of millions!

julia> @time Enzyme.onehot(ψ₀);
  0.134450 seconds (1.92 M allocations: 119.125 MiB)

@profview_allocs points to the implementation inds(tn; set=:open). Such method is being called by copy which calls the Quantum constructor.

[mofeing]

A posible solution might be to use sizehint! together with ninds at the start of the count, and then shrink allocation using sizehint! again.

Even in the abnormal case where we have 1 million indices, the memory overhead of such preemptive allocation would be around 8 MiB... which is totally ok.

julia> Base.summarysize([gensym() for _ in 1:1_000_000]) / 1024^2
7.629432678222656