type stability with ITensors/NDTensors and arrays

ITensor Julia Questions
1

hi, here’s a bit of a tedious question regarding the more inner workings of ITensors and maybe NDTensors, there’s some backstory below but I guess in a nutshell it goes like:

is it normal that a function like

using ITensors, LinearAlgebra
function func1(it::ITensor)
       it_a = array(it)   #or equivalently(?) Array(it, inds(it))
       u,s,vd = svd(it_a)
       end

is type unstable (at least according to @code_warntype/@descend_code_warntype) ? And if so, is there a way to make it stable? (and should I worry about it at all?)

Long story short, I was playing around with tensor decompositions trying to find a way to speedup my code (always a bad idea it seems…) and realized that I could get a reasonable speed gain if I was replacing in one of my functions an ITensor by a julia array for the intermediate steps, using calls like
my_array = Array(my_itensor, inds(my_itensor))

I was a bit surprised by this, but in the end I decided to stop worrying and be happy with it. But then later on when I was inspecting my code for type stability issues I saw a bloodbath there, basically everything starting from the first inds(itensor) was returning me ::Any types, which I’m no super julia expert, but I’m told is really not good for performance.

After some looking around I stumbled upon this page
https://itensor.github.io/ITensors.jl/dev/AdvancedUsageGuide.html#NDTensors-and-ITensors
which to be honest I didn’t fully understand (I’m relatively new to julia’s magic inner workings), but if I get it right it basically says that some extent of type instability can come up when messing around with ITensors, and that maybe working directly with the underlying NDTensors could alleviate that.

So I started digging a bit into NDTensors but got a bit lost there, the documentation seems a bit more scarce than for ITensors (though I probably didn’t look hard enough), I was not sure what the relevant elements/functions for initialising an NDTensor are, I was even a bit puzzled since it seems that the SVD convention is different from the ITensor one (though I saw later a bug report in the archived NDtensor repo that perhaps explains it… ), anyway that got me a bit worried that maybe I was looking at something I wasn’t supposed to play with, so I wonder what’s the official take on this from the developers.

as always, thanks for developing itensors!

2

Accessing the data or indices of the ITensor type is type unstable by design, which allows for a more flexible high level user interface (like being able to change the data type of the ITensor type in-place) at the expense of a small overhead each time the data or indices are accessed.

If you look through the source code of ITensors.jl, you’ll notice that most ITensor functions are written by unwrapping with NDTensors.tensor, performing an operation on the resulting Tensor object (the ITensor object is just a wrapper around a Tensor object which is defined in our internal tensor library NDTensors.jl, and is generally type stable), and then rewrapping in an ITensor. This is called a “function barrier” which is described here in the Julia manual: Performance Tips · The Julia Language. Basically, it minimizes the effects of the type instability of the ITensor type.

DataFrames.jl is a popular package which makes a similar design choice, see this blog post for details: Why DataFrame is not type stable and when it matters | Blog by Bogumił Kamiński.

So in your code above, you might rewrite it as:

using ITensors, LinearAlgebra
func1(it::ITensor) = func1(array(it))
function func1(a::Array)
  # Some operations on an `Array`, for example accessing many elements or `svd`
end

This create a function barrier so that the only performance issue caused by the type instability is a small one for unwrapping the internal data with array, then the inner function func1 defined on Array is as fast as any type stable Julia code will be.

3

Thanks for the info! Though I have to admit I’m still somewhat confused, if I do something like

using ITensors, LinearAlgebra, Test

func2(it::ITensor) = func2(array(it))

function func2(arr::Array) 
    return svd(arr)
end

a = rand(2,2)  # just an array
at = ITensor(a, Index(2), Index(2))  # the itensor on top

Test.@inferred func2(a)  # this seems to work alright

Test.@inferred func2(at)  # this one the type inference doesn't seem to work though

namely for the latter I get
ERROR: return type SVD{Float64, Float64, Matrix{Float64}, Vector{Float64}} does not match inferred return type Any

and I’m mostly worried about the return types of the function and some ::Anys propagating through my code.
Maybe I’m still missing something basic here, I’ll go and read more carefully the performance hints page you linked

4

Matt might have a more detailed answer to give, but the reason here that Julia cannot infer the return type of func2(at) is that array(::ITensor) returns different types depending on the order (number of indices) the ITensor has. So array(at) returns an Array{Float64,2} for your matrix-like ITensor but for an ITensor with three indices, array would return an Array{Float64,3}.

This will nevertheless cost very little overhead in Julia when you use the function-barrier pattern because Julia will just compile a separate implementation of func2(arr::Array) for all the cases like Array{Float63,2},Array{Float63,3}, … and call the necessary one. Each one will be optimized for that number of indices and will individually be fast.

But @inferred is inherently unable to know which of those will actually be called based only on type information since the ITensor type doesn’t store how many indices the ITensor has and it’s only knowable at run time (i.e. dynamically). Another way to say it is that your code really will have a type instability, but it is one that is very cheap and common in Julia and not a kind to worry about.

5

@miles beyond that, Julia also can’t infer the element type of the ITensor since that information also isn’t available to the compiler. If you use a converter Array{Float64,2}(at) (I forget if that is defined) then it should be inferrable.

I think the key assumption is that the bottleneck of your code is the operations that happen inside the function barrier, not outside the function barrier. So even though Julia can’t infer the types in the outer code, it will compile a fast version of the inner function func2(arr::Array), and if that is the bottleneck of your code then there isn’t anything to worry about.

1 Like
6

A common code pattern inside the ITensors.jl and NDTensors.jl libraries is the following:

using ITensors
using NDTensors

scale2(t::ITensor) = itensor(scale2(tensor(t)))

function scale2(t::Tensor)
  t = copy(t)
  scale2!(t)
  return t
end

function scale2!(t::Tensor)
  for i in eachindex(t)
    t[i] *= 2
  end
  return t
end

i, j = Index.((10, 10))
A = randomITensor(i, j)
B = scale2(A)
@show 2A ≈ B

using Test
@inferred scale2(A)

The code here in fact is inferred from the outside because we rewrap the result in an ITensor, which hides some of the type instability.

So the only type unstable code is tensor(t) since the ITensor hides the type of the tensor object, but then it gets wrapped again so from the outside the output of scale2(A) is always just ITensor.

So type instabilities are only seen if a function outputs the data of an ITensor, for example A[i => 1, j => 1] would not be type stable. However, accessing individual elements of an ITensor like that is not recommended in performant code.

7

I see, thanks a lot for all the info!!

so just to satisfy my OCD I checked and indeed if I specify an interface like

fun(it::ITensor) = fun(array(it))

function fun(arr::Array{ComplexF64,2}) 
    return svd(arr)
end

I do get a type stable result!

Though if I do instead

function fun(arr::Array{T,2}) where T<:Union{ComplexF64,Float64}
    return svd(arr)
end

then this is not type stable anymore, and I have no idea why,
but I guess this is me not understanding julia well enough yet, and more importantly I checked and indeed as you said, the final overhead in my code seems quite negligible, so I’m happy with it

thanks a lot!!

8

I’m pretty sure the reason that last one doesn’t pass the @inferred check is that Julia will implement that generic code by compiling two functions:

fun(arr::Array{ComplexF64,2})
fun(arr::Array{Float64,2})

and it cannot tell from the call to fun(::ITensor) which of these two are going to be called, and each of them gives a different return value type.

In the first case, there is only on fun method taking an array, so there Julia can infer the return type because there’s only one possible function that could ever be called.