`DirectLinearSolver` with vectors of length 1

[bmxam]

The DirectLinearSolver is failing one dealing with "vectors" of length 1, see the below MWE:

module MWE
using ForwardDiff
using ImplicitDifferentiation
using LinearAlgebra
using Optim
using StaticArrays

const ε = 0.1 # radial deviation amplitude
const m = 7 # number of lobes

""" For a point `x`, find `θ0` such that the square distance between `x` and `r_star(θ0)` is minimal. """
function forward(x)
    function f(θ)
        _θ = first(θ)
        u = x - (1 + ε * cos(m * _θ)) .* SA[cos(_θ), sin(_θ)]
        return sum(abs2, u)
    end
    θ_init = atan(x[2], x[1])
    results = optimize(f, [θ_init])

    y = Optim.minimizer(results)
    z = nothing

    return y, z
end

""" Here `y` represents `θ0` (see `forward` doc). Gradient, with respect to `θ0` of the squared distance between `x` and `r_star(θ0)`. """
function conditions(x, y, _z)
    x1, x2 = x
    θ = first(y)

    r0 = 1 + ε * cos(m * θ)
    r0_p = -m * ε * sin(m * θ)
    c = cos(θ)
    s = sin(θ)

    ∇₂f = [
        2 * (-r0_p * c + r0 * s) * (x1 - r0 * c) + 2 * (-r0_p * s - r0 * c) * (x2 - r0 * s),
    ]
    return ∇₂f
end

function run()
    p = [1.2, 0.0]

    implicit_iterative = ImplicitFunction(forward, conditions;)
    @show ForwardDiff.jacobian(first ∘ implicit_iterative, p) # OK with `IterativeLinearSolver`

    implicit_direct = ImplicitFunction(
        forward,
        conditions;
        representation = MatrixRepresentation(),
        linear_solver = DirectLinearSolver(),
    )
    @show ForwardDiff.jacobian(first ∘ implicit_direct, p) # Error : `MethodError: no method matching (::ImplicitDifferentiation.DirectLinearSolver)(::Float64, ::Nothing, ::Vector{Float64}, ::Vector{Float64})`
end

run()

end

The cause is here

ImplicitDifferentiation.jl/src/execution.jl

Line 60 in 96c2a72

return factorize(A)

. If A is an array of size (1,1), then the factorize(A) returns a Float (and not a Factorization). To fix it (if we want to fix it), I see at least two options:

• replace the previously mentionned line by a wrap_as_factorization(factorize(A)) that would return a proper object
• implement function (solver::DirectLinearSolver)(A::Float, _Aᵀ, b::AbstractVector, x0::AbstractVector)

I am not sure this should be fixed since it is a very specific case, but I think that at least the behavior of IterativeLinearSolver and DirectLinearSolver should be the same.

In the mean time, as a quick-and-dirty trick, we can return a vector with two variables:

function forward(x)
    function f(θ)
        _θ = first(θ)
        u = x - (1 + ε * cos(m * _θ)) .* SA[cos(_θ), sin(_θ)]
        return sum(abs2, u)
    end
    θ_init = atan(x[2], x[1])
    results = optimize(f, [θ_init])

    y = vcat(Optim.minimizer(results), 0.0)
    z = nothing

    return y, z
end

function conditions(x, y, _z)
    x1, x2 = x
    θ = first(y)

    r0 = 1 + ε * cos(m * θ)
    r0_p = -m * ε * sin(m * θ)
    c = cos(θ)
    s = sin(θ)

    ∇₂f = [
        2 * (-r0_p * c + r0 * s) * (x1 - r0 * c) + 2 * (-r0_p * s - r0 * c) * (x2 - r0 * s),
        last(y),
    ]
    return ∇₂f
end

function run()
    p = [1.2, 0.0]

    implicit_direct = ImplicitFunction(
        forward,
        conditions;
        representation = MatrixRepresentation(),
        linear_solver = DirectLinearSolver(),
    )
    @show ForwardDiff.jacobian(first ∘ implicit_direct, p)
end

[gdalle]

I think just adding Number to the dispatches for the direct solver is good enough:

struct DirectLinearSolver <: AbstractSolver end

function (solver::DirectLinearSolver)(
        A::Union{AbstractMatrix, Factorization, Number}, _Aᵀ, b::AbstractVector, x0::AbstractVector
    )
    return A \ b
end

Again, since you spotted the bug, feel free to open a PR with this and a boiled down version of your example as a test, otherwise I'll do it in the next couple of days! I don't want to exploit you, just want to give you the opportunity to contribute if it matters to you.
Thanks again