Simple Reshaping in Zygote Throwing MethodError

Specific Domains Machine Learning
2

The problem is that there’s no AD rule for reduce(vcat, ::AbstractMatrix) and res' is a matrix. This can be worked around by taking the adjoint of the inner arrays only:

function ffunc(x::Vector{Float64})
    # first reshape
    xMat = reshape(x, (:, 2))
    # apply function
    res = [
        (xMat[:,1] .+ xMat[:,2])',
        (xMat[:,1] .- xMat[:,2])'
    ] # return value is a vector of two vectors

    return vec(reduce(vcat, res))
end

I’ve taken the liberty of using vec instead of the outer vcat because it’s more AD-friendly and faster.

If there are always a fixed number of arrays in res, however, you could consider making it not an array:

function ffunc(x::Vector{Float64})
    # first reshape
    xMat = reshape(x, (:, 2))
    # apply function
    res = (
        (xMat[:,1] .+ xMat[:,2])',
        (xMat[:,1] .- xMat[:,2])'
    ) # return value is a tuple of two vectors

    return vec(vcat(res...))
end

Splatting generally works fine with tuples under Zygote.