JuliaDiff
Nested jacobian! calls
Hi!
I'm having some troubles using nested inplace jacobian:
using ForwardDiff
using LinearAlgebra
#some data
x0 = randn(50000);
p0 = randn(2);
#pre allocate outputs
out = similar(x0);
outjac = Array{Float64}(undef, 50000, 2);
outjac2 = Array{Float64}(undef,100000, 2)
#'regular' and in place function
f(p) = @. p[1] * exp(-x0 * p[2]);
f!(out, p) = @. out = p[1] * exp(-x0 * p[2]);
ForwardDiff.jacobian!(outjac2,(outjacin,p) -> (outjacin .= ForwardDiff.jacobian(f, p)), outjac, p0) #works like a charm
ForwardDiff.jacobian!(outjac2,(outjacin,p) -> ForwardDiff.jacobian!(outjacin, f!, out, p), outjac, p0); # breaks
It throws MethodError: no method matching Float64(::ForwardDiff.Dual{ForwardDiff.Tag{getfield(Main, Symbol("##77#78")),Float64},Float64,2})
Both inner functions work and produce the same result:
g1 = (outjacin,p) -> ForwardDiff.jacobian!(outjacin, f!, out, p)
g1(outjac,p0)
g2 = (outjacin,p) -> (outjacin .= ForwardDiff.jacobian(f, p))
g2(outjac,p0)
Thanks for the help!
A similar problem:
jac1 = p -> ForwardDiff.jacobian(f, p)
jac2 = p -> ForwardDiff.jacobian(f!,out, p)
jac1(p0) == jac2(p0)
> true
ForwardDiff.jacobian(jac1, p0) #works
ForwardDiff.jacobian(jac2, p0) #crashes
outputs
MethodError: no method matching Float64(::ForwardDiff.Dual{ForwardDiff.Tag{getfield(Main, Symbol("##130#131")),Float64},Float64,2})
https://discourse.julialang.org/t/nested-forwarddiff-jacobian-calls-with-inplace-function/21232/2
https://discourse.julialang.org/t/optim-gradient-and-hessian-of-constraints/31493/3 here i wrote a code that does a jacobian of a jacobian of an inplace function for use in Optim, maybe that can help
I have recently come across this and here is my take on the issue, very much inspired by how hessian! is implemented.
We take care to define duals for the nested jacobian calls.
Due to my use-case, I have coded vector_hessian! to mutate a DiffResult.
As all target arrays are contained in result already, I deviate from the standard syntax a bit: To indicate an in-place function with signature f!(y, x), set the kwarg f_is_in_place=true.
Additionally, I'd like my Hessians in “tensor-format”, so there are some non-allocating transformations in case that ndims(DiffResults.hessian(result))==3 :)
import ForwardDiff.DiffResults
import ForwardDiff: jacobian, jacobian!, JacobianConfig, DiffResult, MutableDiffResult,
Dual, Chunk, Tag, value, checktag
struct VectorHessianConfig{T, V, N, DI, DO}
inner_config::JacobianConfig{T, Dual{T, V, N}, N, DI}
outer_config::JacobianConfig{T, V, N, DO}
ydual::Vector{Dual{T,V,N}}
end
function VectorHessianConfig(
f::F, result::DiffResult, x::AbstractArray{V}, chunk::Chunk = Chunk(x), tag = Tag(f, V)
) where {F, V}
Dy = DiffResults.jacobian(result)
nout = size(Dy, 1)
outer_config = JacobianConfig((f, jacobian), Dy, x, chunk, tag)
xdual = outer_config.duals[2]
ydual = zeros(eltype(xdual), nout)
inner_config = JacobianConfig(f, ydual, xdual, chunk, tag)
return VectorHessianConfig(inner_config, outer_config, ydual)
end
mutable struct InnerJacobianForHess!{R,C,F}
result::R
cfg::C
f::F
f_is_in_place::Bool
end
function (g::InnerJacobianForHess!)(Dydual, xdual)
inner_result = DiffResult(g.cfg.ydual, Dydual)
if g.f_is_in_place
jacobian!(inner_result, g.f, g.cfg.ydual, xdual, g.cfg.inner_config, Val(false))
else
jacobian!(inner_result, g.f, xdual, g.cfg.inner_config, Val(false))
end
# jacobian!(inner_result, g.f, xdual, g.cfg.inner_config, Val{false}(); f_is_in_place=g.f_is_in_place)
g.result = DiffResults.value!(g.result, value.(DiffResults.value(inner_result)))
return nothing
end
function vector_hessian!(
result::MutableDiffResult{O},
f::F,
x::AbstractArray,
cfg::VectorHessianConfig{T} = VectorHessianConfig(f, result, x), ::Val{CHK}=Val{true}();
f_is_in_place::Bool=false,
) where {O,F,T,CHK}
@assert O >= 2 "Target `result` must contain arrays for second-order derivatives."
CHK && checktag(T, f, x)
∇f! = InnerJacobianForHess!(result, cfg, f, f_is_in_place)
Dy = DiffResults.jacobian(result)
nout, nvars = size(Dy)
H = DiffResults.hessian(result)
h = if ndims(H) == 3
#=
Assume, we want to store the Hessians in a 3D array `H`, such that `H[:,:,k]`
is the `k`-th Hessian matrix, i.e., `H[i,j,k]` is ∂/∂ᵢ∂ⱼ Hₖ.
Given the Jacobian
```
J = [
∂₁f₁ ∂₂f₁ … ∂ₙf₁ ;
∂₁f₂ ∂₂f₂ … ∂ₙf₂ ;
⋮ ⋮ ⋱ ⋮
∂₁fₘ ∂₂fₘ … ∂ₙfₘ
]
```
the function `ForwardDiff.jacobian` interprets the colums as seperate outputs
and stacks their respective Jacobians vertically.
Consider `J[:,i] = [∂ᵢf₁, …, ∂ᵢfₘ]`.
Its Jacobian is
```
∇(J[:,i]) = [
(∇∂ᵢf₁)ᵀ;
(∇∂ᵢf₂)ᵀ;
⋮
(∇∂ᵢfₘ)ᵀ;
] = [
∂₁∂ᵢf₁ ∂₂∂ᵢf₁ … ∂ₙ∂ᵢf₁ ;
∂₁∂ᵢf₂ ∂₂∂ᵢf₂ … ∂ₙ∂ᵢf₂ ;
⋮ ⋮ ⋱ ⋮
∂₁∂ᵢfₘ ∂₂∂ᵢfₘ … ∂ₙ∂ᵢfₘ ;
]
```
Thus, the `i`-th block of the stacked Hessian contains as its rows the
`i`-th columns of all function Hessians.
To get `H` from `∇J`, we can do
`H = permutedims(reshape(transpose(reshape(∇J, m, n*n)), n, n, m), (2,1,3))`
But if `H` is provided as a pre-allocated buffer, we need the inverse operation:
`∇J = reshape(transpose(reshape(permutedims(H, (2,1,3)), n*n, m)), m*n, n)`
In case of symmetric Hessians, we could forgo the `permutedims` operation...
=#
reshape(transpose(reshape(PermutedDimsArray(H, (2,1,3)), nvars*nvars, nout)), nout*nvars, nvars)
else
H
end
jacobian!(h, ∇f!, Dy, x, cfg.outer_config, Val{false}())
return nothing
end
https://github.com/JuliaDiff/ForwardDiff.jl/issues/360 might be related?
ok, I actually had to change the signatures a bit, to make it work in a fresh REPL:
import ForwardDiff.DiffResults
import ForwardDiff: jacobian, jacobian!, JacobianConfig, DiffResult, MutableDiffResult,
Dual, Chunk, Tag, value, checktag
struct VectorHessianConfig{T, V, N, DI, DO}
inner_config::JacobianConfig{T, Dual{T, V, N}, N, DI}
outer_config::JacobianConfig{T, V, N, DO}
ydual::Vector{Dual{T,V,N}}
end
function VectorHessianConfig(
f::F, result::DiffResult, x::AbstractArray{V}, f_is_in_place::Bool=false, chunk::Chunk=Chunk(x), tag=Tag(f, V)
) where {F, V}
Dy = DiffResults.jacobian(result)
outer_config = JacobianConfig((f, jacobian), Dy, x, chunk, tag)
xdual = outer_config.duals[2]
nout = size(Dy, 1)
ydual = zeros(eltype(xdual), nout)
inner_config = if f_is_in_place
JacobianConfig(f, ydual, xdual, chunk, tag)
else
JacobianConfig(f, xdual, chunk, tag)
end
return VectorHessianConfig(inner_config, outer_config, ydual)
end
mutable struct InnerJacobianForHess!{R,C,F}
result::R
cfg::C
f::F
f_is_in_place::Bool
end
function (g::InnerJacobianForHess!)(Dydual, xdual)
inner_result = DiffResult(g.cfg.ydual, Dydual)
if g.f_is_in_place
jacobian!(inner_result, g.f, g.cfg.ydual, xdual, g.cfg.inner_config, Val(false))
else
jacobian!(inner_result, g.f, xdual, g.cfg.inner_config, Val(false))
end
g.result = DiffResults.value!(g.result, value.(DiffResults.value(inner_result)))
return nothing
end
function vector_hessian!(
result::MutableDiffResult{O},
f::F,
x::AbstractArray,
f_is_in_place::Bool=false,
cfg::VectorHessianConfig{T}=VectorHessianConfig(f, result, x, f_is_in_place),
::Val{CHK}=Val{true}();
) where {O,F,T,CHK}
@assert O >= 2 "Target `result` must contain arrays for second-order derivatives."
CHK && checktag(T, f, x)
∇f! = InnerJacobianForHess!(result, cfg, f, f_is_in_place)
Dy = DiffResults.jacobian(result)
nout, nvars = size(Dy)
H = DiffResults.hessian(result)
h = if ndims(H) == 3
reshape(transpose(reshape(PermutedDimsArray(H, (2,1,3)), nvars*nvars, nout)), nout*nvars, nvars)
else
H
end
jacobian!(h, ∇f!, Dy, x, cfg.outer_config, Val{false}())
return nothing
end