Unable to infer ranges for indices `i`, in deeply nested `r[k][j][i]`

[tpapp]

Consider the MWE:

using TensorCast # 0.4.8

r = fill(fill(ones(3), 5), 7)

function f(r)
    @cast _[i, j ⊗ k] := r[k].a[j][i]
end

Ideally f(r) should give a 3 x 35 matrix, but I get

ERROR: LoadError: unable to infer ranges for indices i
    @cast _[i, j ⊗ k] := ((r[k]).a[j])[i]  
    @ Main REPL[22]:2
Stacktrace:
 [1] sizeinfer(store::@NamedTuple{dict::Dict{Any, Any}, assert::Vector{Any}, mustassert::Vector{Any}, seen::Vector{Any}, need::Vector{Any}, top::Vector{Any}, main::Vector{Any}}, call::TensorCast.CallInfo)
   @ TensorCast ~/.julia/packages/TensorCast/fctFg/src/macro.jl:1143
 [2] findsizes(store::@NamedTuple{dict::Dict{Any, Any}, assert::Vector{Any}, mustassert::Vector{Any}, seen::Vector{Any}, need::Vector{Any}, top::Vector{Any}, main::Vector{Any}}, call::TensorCast.CallInfo)
   @ TensorCast ~/.julia/packages/TensorCast/fctFg/src/macro.jl:1092
 [3] _macro(exone::Expr, extwo::Nothing, exthree::Nothing; call::TensorCast.CallInfo, dict::Dict{Any, Any})
   @ TensorCast ~/.julia/packages/TensorCast/fctFg/src/macro.jl:211
 [4] _macro
   @ ~/.julia/packages/TensorCast/fctFg/src/macro.jl:154 [inlined]
 [5] var"@cast"(__source__::LineNumberNode, __module__::Module, exs::Vararg{Any})
   @ TensorCast ~/.julia/packages/TensorCast/fctFg/src/macro.jl:74
in expression starting at REPL[22]:2

Is it possible to make this work somehow?

[tpapp]

Note that breaking it down works:

function g(a)
    @cast _[i, j] := a[j][i]
end

function h(r)
    z = map(g, r)
    @cast _[i, j ⊗ k] := z[k][i, j]
end

h(r)

[mcabbott]

Can you clarify what a is, and what the dot in ((r[k]).a[j])[i] means? r[1] isa Vector{Vector{Float64}} which doesn't seem to have useful fields.

From your second message, I think you want this:

julia> r = [[rand(1:99, 3) for _ in 1:5] for _ in 1:7];  # ensure order of j ⊗ k matters

julia> @show size(r) size(r[1]) size(r[1][1]);  # same sizes
size(r) = (7,)
size(r[1]) = (5,)
size((r[1])[1]) = (3,)

julia> h(r) ≈ @cast _[i, j ⊗ k] := r[k][j][i]  i in 1:3
true

julia> h(r) ≈ @cast _[i, j ⊗ k] := r[j][k][i]  i in 1:3
false

julia> h(r) ≈ @cast _[i, j ⊗ k] := r[j][k][i]
ERROR: LoadError: unable to infer ranges for indices i
    @cast _[i, j ⊗ k] := ((r[j])[k])[i]  
    @ Main REPL[23]:1
Stacktrace:
 [1] sizeinfer(store::@NamedTuple{…}, call::TensorCast.CallInfo)
   @ TensorCast ~/.julia/packages/TensorCast/fctFg/src/macro.jl:1143
 [2] findsizes(store::@NamedTuple{…}, call::TensorCast.CallInfo)
   @ TensorCast ~/.julia/packages/TensorCast/fctFg/src/macro.jl:1092

I'm a little surprised that it can't infer the range of i implied by r[j][k][i], as I thought most things were fully recursive. So most likely that counts as a bug. But I'd have to remind myself how findsizes works to be sure.

[tpapp]

Can you clarify what a is, and what the dot in ((r[k]).a[j])[i] means?

Sorry, the original code had an extra level of indirection. Assume

r = fill((a = fill(ones(3), 5), ), 7)

for the first example. I lost track when I tried to simplify it.

[mcabbott]

OK. This works the same way, in that the code to stack it seems happy with this level of recursion, but the code to infer the ranges of indices seems unable to recurse to the innermost array:

julia> r2 = fill((a = fill(ones(3), 5), ), 7);

julia> @cast _[i, j ⊗ k] := r2[j].a[k][i]  i in 1:3
3×35 reshape(transmute(lazystack(lazystack(::Vector{Vector{Vector{Float64}}})), (1, 3, 2)), 3, 35) with eltype Float64:
 1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  …  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0
 1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0     1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0
 1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0     1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0

julia> @cast _[i, j ⊗ k] |= r2[j].a[k][i]  i in 1:3
3×35 Matrix{Float64}:
 1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  …  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0
 1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0     1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0
 1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0     1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0

julia> @cast _[i, j ⊗ k] := r2[j].a[k][i]
ERROR: LoadError: unable to infer ranges for indices i
    @cast _[i, j ⊗ k] := ((r2[j]).a[k])[i]  
    @ Main REPL[67]:1
Stacktrace:
 [1] sizeinfer(store::@NamedTuple{…}, call::TensorCast.CallInfo)
   @ TensorCast ~/.julia/packages/TensorCast/fctFg/src/macro.jl:1143
 [2] findsizes(store::@NamedTuple{…}, call::TensorCast.CallInfo)
   @ TensorCast ~/.julia/packages/TensorCast/fctFg/src/macro.jl:1092

As you say, one work-around is to do this in two stages. The arrays being made here are mostly lazy views, so this may not cost more allocations. (The .a is a broadcasted getproperty which is eager, but done once regardless.)

Another possible work-around is to add some other expression by which it can infer the size. For example:

julia> @cast _[i, j ⊗ k] := r2[j].a[k][i] + 0 * first(first(r2).a)[i]
3×35 Matrix{Float64}:
 1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  …  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0
 1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0     1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0
 1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0     1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0

julia> @pretty @cast _[i, j ⊗ k] := r2[j].a[k][i] + 0 * first(first(r2).a)[i]
begin
    local var"first((first(r2)).a)_val" = first((first(r2)).a)
    @boundscheck r2 isa Tuple || (ndims(r2) == 1 || throw(ArgumentError("expected a vector or tuple r2[j]")))
    @boundscheck var"first((first(r2)).a)_val" isa Tuple || (ndims(var"first((first(r2)).a)_val") == 1 || throw(ArgumentError("expected a vector or tuple first((first(r2)).a)_val[i]")))
    local (ax_i, ax_j, ax_k) = (axes(var"first((first(r2)).a)_val", 1), axes(r2, 1), axes((first(r2)).a, 1))
    local fly = @__dot__(getproperty(r2, :a))
    local penguin = transmute(lazystack(lazystack(fly)), Val((1, 3, 2)))
    locust = reshape(@__dot__(penguin + 0 * var"first((first(r2)).a)_val"), (ax_i, star(ax_j, ax_k)))
end