Type Inference for CartesianIndices(array) fail when the dimension of the array is too large

I do a lot of work with high-dimensional arrays, where the dimension of my array is often 20+. I noticed that CartesianIndices fail when the dimension is too large, causing later type inferences to also fail. Here is a minimal example with Julia 1.2.0 where that behavior can be observes.

using BenchmarkTools

function f(array)
    a = CartesianIndices(array)
    return a
end

#input is 15-dimensional array
small = rand(Float64,(2*ones(Int64, 15))...)
#input is 16-dimensional array
big = rand(Float64,(2*ones(Int64,16))...)

@code_warntype f(small)
@code_warntype f(big)

The output is given by the following.

Variables
  #self#::Core.Compiler.Const(f, false)
  array::Array{Float64,15}
  a::CartesianIndices{15,NTuple{15,Base.OneTo{Int64}}}     

Body::CartesianIndices{15,NTuple{15,Base.OneTo{Int64}}}    
1 ─     (a = Main.CartesianIndices(array))
└──     return a

Variables
  #self#::Core.Compiler.Const(f, false)
  array::Array{Float64,16}
  a::Any

Body::Any
1 ─     (a = Main.CartesianIndices(array))
└──     return a

As you can see, type inference failed for the bigger array.

Is there a fix for this, or do we just need to avoid this failure to propage with e.g. function barriers?

Things go wrong here:

julia> @code_warntype axes(a)
Variables
  #self#::Core.Compiler.Const(axes, false)
  A::Array{Float64,16}

Body::Tuple
1 ─      nothing
β”‚   %2 = Base.size(A)::NTuple{16,Int64}
β”‚   %3 = Base.map(Base.OneTo, %2)::Tuple
└──      return %3

map for tuples doesn’t infer after N = 15. This limit could be increased in Base, but that would mean increased code size and longer compilation times for certain programs. There are also other ways to maintain inferability (switching from map to a generated function or using ntuple with a Val argument, for example). That may be warranted in this case. I’m guessing there will be many functions like this that will stop inferring for very high-dimensional arrays though.

Just curious, what use case do you have for such high-dimensional arrays?

1 Like

Thanks for the response. That makes sense. I guess for now I need to use functions barrier or avoid using CartesianIndices.

The use case I have in mind is tensor network for quantum physics. There, each index can represent different physical degrees of freedom, and depending on the number of degrees of freedom, it can become rather high-dimensional. i can of course always use linear indexing instead, but I believe using high-dimensional array make the code more transparent to physicist readers.

If you want to this to be fixed long-term, you could open a PR against base with something like

function _axes(sz::NTuple{N, Int}) where N
    ntuple(i -> Base.OneTo(sz[i]), Val(N))
end

Base.axes(A) = _axes(size(A))

to replace

You can already run the code above in the REPL to do an ad-hoc replacement of the axes definition, which indeed fixes the type inference issue in CartesianIndices. But that’s type piracy, so a base PR would be the way to go.