Error in get_ancestor_indices when calling with type CartesianIndex

[Todorbsc]

When doing Reactant compilation of a function, I ran into a bad "unpacking" of indices inside get_ancestor_indices function for the case of a CartesianIndex{N} type. The stacktrace is shown below, where I added @show indices N to print the issue:

julia> function f(params)
           @allowscalar U = Quantum(convert(Circuit, dispatch(yaocirc, params)))
           return @allowscalar U'
       end
f (generic function with 1 method)

julia> fR = @compile f(params′)
indices = (CartesianIndex(1, 1),)
N = 2
ERROR: AssertionError: Expected 2 indices, got 1
Stacktrace:
  [1] get_ancestor_indices(x::Base.ReshapedArray{…}, indices::CartesianIndex{…})
    @ Reactant.TracedUtils ~/.julia/packages/Reactant/gCoju/src/TracedUtils.jl:67
  [2] getindex
    @ ~/.julia/packages/Reactant/gCoju/src/TracedRArray.jl:213 [inlined]
  [3] getindex(none::Base.ReshapedArray{…}, none::Tuple{…})
    @ Reactant ./<missing>:0
  [4] getindex
    @ ~/.julia/packages/Reactant/gCoju/src/TracedRArray.jl:213 [inlined]
  [5] call_with_reactant(::Reactant.MustThrowError, ::typeof(getindex), ::Base.ReshapedArray{…}, ::CartesianIndex{…})
    @ Reactant ~/.julia/packages/Reactant/gCoju/src/utils.jl:0
  [6] getindex
    @ ~/.julia/packages/Tenet/qEDim/src/Tensor.jl:172 [inlined]
  [7] _broadcast_getindex
    @ ./broadcast.jl:639 [inlined]
  [8] _getindex
    @ ./broadcast.jl:670 [inlined]
  [9] _broadcast_getindex
    @ ./broadcast.jl:645 [inlined]
 [10] getindex
    @ ./broadcast.jl:605 [inlined]
 [11] macro expansion
    @ ./broadcast.jl:968 [inlined]
 [12] macro expansion
    @ ./simdloop.jl:77 [inlined]
 [13] copyto!
    @ ./broadcast.jl:967 [inlined]
 [14] copyto!(none::Tensor{…}, none::Base.Broadcast.Broadcasted{…})
    @ Reactant ./<missing>:0
 [15] getproperty
    @ ./Base.jl:49 [inlined]
 [16] parent
    @ ~/.julia/packages/Tenet/qEDim/src/Tensor.jl:140 [inlined]
 [17] axes
    @ ~/.julia/packages/Tenet/qEDim/src/Tensor.jl:210 [inlined]
 [18] copyto!
    @ ./broadcast.jl:956 [inlined]
 [19] call_with_reactant(::typeof(copyto!), ::Tensor{…}, ::Base.Broadcast.Broadcasted{…})
    @ Reactant ~/.julia/packages/Reactant/gCoju/src/utils.jl:0
 [20] copyto!
    @ ./broadcast.jl:920 [inlined]
 [21] copy
    @ ./broadcast.jl:892 [inlined]
 [22] materialize
    @ ./broadcast.jl:867 [inlined]
 [23] broadcast_preserving_zero_d
    @ ./broadcast.jl:856 [inlined]
 [24] conj
    @ ./abstractarraymath.jl:146 [inlined]
 [25] conj(none::Tensor{Reactant.TracedRNumber{…}, 2, Base.ReshapedArray{…}})
    @ Reactant ./<missing>:0
 [26] broadcasted
    @ ./broadcast.jl:1328 [inlined]
 [27] broadcasted
    @ ./broadcast.jl:1320 [inlined]
 [28] broadcast_preserving_zero_d
    @ ./broadcast.jl:855 [inlined]
 [29] conj
    @ ./abstractarraymath.jl:146 [inlined]
 [30] call_with_reactant(::typeof(conj), ::Tensor{Reactant.TracedRNumber{…}, 2, Base.ReshapedArray{…}})
    @ Reactant ~/.julia/packages/Reactant/gCoju/src/utils.jl:0
 [31] _broadcast_getindex_evalf
    @ ./broadcast.jl:673 [inlined]
 [32] _broadcast_getindex
    @ ./broadcast.jl:646 [inlined]
 [33] _getindex
    @ ./broadcast.jl:670 [inlined]
 [34] _getindex
    @ ./broadcast.jl:669 [inlined]
 [35] _broadcast_getindex
    @ ./broadcast.jl:645 [inlined]
 [36] getindex
    @ ./broadcast.jl:605 [inlined]
 [37] copy
    @ ./broadcast.jl:906 [inlined]
 [38] copy(none::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{…}, Tuple{…}, Type{…}, Tuple{…}})
    @ Reactant ./<missing>:0
 [39] call_with_reactant(::typeof(copy), ::Base.Broadcast.Broadcasted{…})
    @ Reactant ~/.julia/packages/Reactant/gCoju/src/utils.jl:423
 [40] materialize
    @ ./broadcast.jl:867 [inlined]
 [41] conj
    @ ~/.julia/packages/Tenet/qEDim/src/TensorNetwork.jl:124 [inlined]
 [42] conj(none::Quantum)
    @ Reactant ./<missing>:0
 [43] getproperty
    @ ./Base.jl:49 [inlined]
 [44] TensorNetwork
    @ ~/.julia/packages/Tenet/qEDim/src/Quantum.jl:79 [inlined]
 [45] copy
    @ ~/.julia/packages/Tenet/qEDim/src/Quantum.jl:81 [inlined]
 [46] conj
    @ ~/.julia/packages/Tenet/qEDim/src/TensorNetwork.jl:123 [inlined]
 [47] call_with_reactant(::typeof(conj), ::Quantum)
    @ Reactant ~/.julia/packages/Reactant/gCoju/src/utils.jl:0
 [48] adjoint
    @ ~/.julia/packages/Tenet/qEDim/src/Quantum.jl:361 [inlined]
 [49] macro expansion
    @ ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:206 [inlined]
 [50] f
    @ ./REPL[6]:3 [inlined]
 [51] f(none::Reactant.TracedRArray{Float64, 1})
    @ Reactant ./<missing>:0
 [52] current_task
    @ ./task.jl:152 [inlined]
 [53] task_local_storage
    @ ./task.jl:280 [inlined]
 [54] macro expansion
    @ ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:204 [inlined]
 [55] f
    @ ./REPL[6]:2 [inlined]
 [56] call_with_reactant(::typeof(f), ::Reactant.TracedRArray{Float64, 1})
    @ Reactant ~/.julia/packages/Reactant/gCoju/src/utils.jl:0
 [57] make_mlir_fn(f::Function, args::Tuple{…}, kwargs::Tuple{}, name::String, concretein::Bool; toscalar::Bool, return_dialect::Symbol, args_in_result::Symbol, construct_function_without_args::Bool, do_transpose::Bool)
    @ Reactant.TracedUtils ~/.julia/packages/Reactant/gCoju/src/TracedUtils.jl:332
 [58] make_mlir_fn
    @ ~/.julia/packages/Reactant/gCoju/src/TracedUtils.jl:152 [inlined]
 [59] compile_mlir!(mod::Reactant.MLIR.IR.Module, f::Function, args::Tuple{…}, callcache::Dict{…}; optimize::Bool, no_nan::Bool, backend::String)
    @ Reactant.Compiler ~/.julia/packages/Reactant/gCoju/src/Compiler.jl:589
 [60] compile_mlir! (repeats 2 times)
    @ ~/.julia/packages/Reactant/gCoju/src/Compiler.jl:564 [inlined]
 [61] compile_xla(f::Function, args::Tuple{ConcreteRArray{…}}; client::Nothing, kwargs::@Kwargs{no_nan::Bool, optimize::Bool})
    @ Reactant.Compiler ~/.julia/packages/Reactant/gCoju/src/Compiler.jl:1360
 [62] compile_xla
    @ ~/.julia/packages/Reactant/gCoju/src/Compiler.jl:1339 [inlined]
 [63] compile(f::Function, args::Tuple{…}; sync::Bool, kwargs::@Kwargs{…})
    @ Reactant.Compiler ~/.julia/packages/Reactant/gCoju/src/Compiler.jl:1399

CC @mofeing

[Todorbsc]

Since I don't have permissions to push a branch to the repo and open a PR, I write here the workaround that I did to solve this issue.
Basicly, I added to src/TracedUtils.jl the next function:

get_ancestor_indices(
    x::WrappedReshapedArray{TracedRNumber{T},N,TracedRArray{T,M}}, indices::CartesianIndex{N}
) where {T,N,M} = get_ancestor_indices(x, indices.I...)

[giordano]

You can always create a pull request from a fork.