Request help in understanding why certain operations in a generated function lead to a `MethodError` but are fine outside the generated function


#1

Hi all,

I’m working on adding more specializations to kron for StaticArray and StaticVector types as part of this pull request: https://github.com/JuliaArrays/StaticArrays.jl/pull/407

A suggestion was to try to create the output matrix of the Kronecker product in one go, in contrast to the manner in which the solution is currently implemented. I intend to do this with a comprehension. For example,

using StaticArrays, Base.Test

A = @SMatrix [1 2; 3 4]
B = @SMatrix [5 6; 7 8]
P = [1 2; 3 4]
Q = [5 6; 7 8]
SA = size(A)
SB = size(B)

# The Kronecker product as a comprehension
M = [ A[ia,ja] * B[ib,jb] for ib in 1:SB[1],ia in 1:SA[1],jb in 1:SB[2],ja in 1:SA[2]]
N = similar_type(M,Size((4,4)))(M)
@test N == kron(P,Q)

Now I wish to put this logic into a generated function. I am not entirely sure what I am doing with this meta programming, so my first attempt was:

const _length_limit = Length(200)
mykron1(a::StaticMatrix, b::StaticMatrix) = _mykron1(_length_limit, Size(a), Size(b), a, b)
@generated function _mykron1(::Length{length_limit}, ::Size{SA}, ::Size{SB}, a, b) where {length_limit,SA,SB}
    outsize = SA .* SB
    if prod(outsize) > length_limit
        return :( SizedMatrix{$(outsize[1]),$(outsize[2])}( kron(drop_sdims(a), drop_sdims(b)) ) )
    end
    M = [ a[ia,ja] * b[ib,jb] for ib in 1:SB[1],ia in 1:SA[1],jb in 1:SB[2], ja in 1:SA[2] ]
    return quote
        @_inline_meta
        @inbounds return similar_type($(M),Size($(outsize)))($(M))  
    end
end

This does not work. It gives the following error:

ERROR: MethodError: Cannot `convert` an object of type Int64 to an object of type StaticArrays.SArray{Tuple{2,2},Int64,2,4}
This may have arisen from a call to the constructor StaticArrays.SArray{Tuple{2,2},Int64,2,4}(...),
since type constructors fall back to convert methods.
Stacktrace:
 [1] getindex(::Type{StaticArrays.SArray{Tuple{2,2},Int64,2,4}}, ::Int64, ::Int64) at ./array.jl:208
 [2] (::##5#6{DataType,DataType})(::NTuple{4,Int64}) at ./<missing>:0
 [3] collect(::Base.Generator{Base.Iterators.Prod{UnitRange{Int64},Base.Iterators.Prod{UnitRange{Int64},Base.Iterators.Prod2{UnitRange{Int64},UnitRange{Int64}}}},##5#6{DataType,DataType}}) at ./array.jl:475
 [4] _mykron1(...) at /home/zygmunt/.julia/v0.6/StaticArrays/test/debug2.jl:24
 [5] mykron1(::StaticArrays.SArray{Tuple{2,2},Int64,2,4}, ::StaticArrays.SArray{Tuple{2,2},Int64,2,4}) at /home/zygmunt/.julia/v0.6/StaticArrays/test/debug2.jl:18

Andy Ferris suggested the following amendment to the return statement.

mykron2(a::StaticMatrix, b::StaticMatrix) = _mykron2(_length_limit, Size(a), Size(b), a, b)
@generated function _mykron2(::Length{length_limit}, ::Size{SA}, ::Size{SB}, a, b) where {length_limit,SA,SB}
    outsize = SA .* SB
    if prod(outsize) > length_limit
        return :( SizedMatrix{$(outsize[1]),$(outsize[2])}( kron(drop_sdims(a), drop_sdims(b)) ) )
    end
    M = [ a[ia,ja] * b[ib,jb] for ib in 1:SB[1], ia in 1:SA[1], jb in 1:SB[2], ja in 1:SA[2] ]
    return quote
        @_inline_meta
        @inbounds return similar_type($a, Size($(outsize)))(tuple($(M...)))
    end
end

but this also does not work and gives the same error. After searching the Discourse forum I wonder if my problem is related to:

Not sure how I should be nesting another generated function in the comprehension in this instance? I’d be grateful if someone could suggest a fix for the aforementioned code.


#2

Your error tells you that the comprehension outside the quote can’t access the elements of a and b.

@generated function _mykron2(::Length{length_limit}, ::Size{SA}, ::Size{SB}, a, b) where {length_limit,SA,SB}
    outsize = SA .* SB
    if prod(outsize) > length_limit
        return :( SizedMatrix{$(outsize[1]),$(outsize[2])}( kron(drop_sdims(a), drop_sdims(b)) ) )
    end
    M = [ :(a[$ia,$ja] * b[$ib,$jb]) for ib in 1:SB[1], ia in 1:SA[1], jb in 1:SB[2], ja in 1:SA[2] ]
    return quote   
        @inbounds return similar_type(a, Size($(outsize)))(tuple($(M...)))
    end
end
julia> mykron2(A, B)
4×4 StaticArrays.SArray{Tuple{4,4},Int64,2,16}:
  5   6  10  12
  7   8  14  16
 15  18  20  24
 21  24  28  32

(Unsure where the @_inline_meta comes from so I removed it in my example)