Lazy Jacobian multiplication

[gdalle]

I had to write this code recently for ImplicitDifferentiation.jl, and @mohamed82008 suggested it might feel at home here. Any opinions?

"""
   LazyJacobianMul!{M,N}

Callable structure wrapping a lazy Jacobian operator with `N`-dimensional inputs into an in-place multiplication for vectors.
# Fields
- `J::M`: the lazy Jacobian of the function
- `input_size::NTuple{N,Int}`: the array size of the function input
"""
struct LazyJacobianMul!{M<:LazyJacobian,N}
    J::M
    input_size::NTuple{N,Int}
end

"""
    LazyJacobianTransposeMul!{M,N}

Callable structure wrapping a lazy Jacobian operator with `N`-dimensional outputs into an in-place multiplication for vectors.
# Fields
- `J::M`: the lazy Jacobian of the function
- `output_size::NTuple{N,Int}`: the array size of the function output
"""
struct LazyJacobianTransposeMul!{M<:LazyJacobian,N}
    J::M
    output_size::NTuple{N,Int}
end

function (ljm::LazyJacobianMul!)(res::Vector, δinput_vec::Vector)
    (; J, input_size) = ljm
    δinput = reshape(δinput_vec, input_size)
    δoutput = only(J * δinput)
    return res .= vec(δoutput)
end

function (ljtm::LazyJacobianTransposeMul!)(res::Vector, δoutput_vec::Vector)
    (; J, output_size) = ljtm
    δoutput = reshape(δoutput_vec, output_size)
    δinput = only(δoutput' * J)
    return res .= vec(δinput)
end

[devmotion]

Can you explain the motivation for it? I assumed that typically the main idea is to not construct the Jacobian at all but directly compute Jacobian-vector products. But here you just store the Jacobian?

[devmotion]

Oh wait, it's actually a lazy operator that is stored? Then you could just define mul!, couldn't you?

[gdalle]

The motivation was to use this lazy Jacobian as a LinearOperator (in the sense of LinearOperators.jl) in our package ImplicitDifferentiation.jl. It's easier to define mul! on a struct, and I needed to know the array order statically in order to write type-stable code

[gdalle]

But I eventually gave up because of #83

[oxinabox]

Isn't this what our lazy_jacobian function already returns?

[gdalle]

Maybe but I needed the additional flattening step and the existence of a mul! method because of the Krylov solver we were using.

[gdalle]

Closing in favor of #114