# Nested ForwardDiff.jacobian calls with inplace function

**URL:** https://discourse.julialang.org/t/nested-forwarddiff-jacobian-calls-with-inplace-function/21232
**Category:** General Usage
**Created:** [February 26, 2019, 8:27pm UTC](https://discourse.julialang.org/t/nested-forwarddiff-jacobian-calls-with-inplace-function/21232 "2019-02-26T20:27:34Z")
**Posts on this page:** 1
**Showing post:** 3

<div class="post-metadata">

### Author: ![tkoolen](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/tkoolen/32/1603_2.png) [@tkoolen](https://discourse.julialang.org/u/tkoolen)
#### Post date: [February 27, 2019, 2:29am UTC](https://discourse.julialang.org/t/nested-forwarddiff-jacobian-calls-with-inplace-function/21232/3 "2019-02-27T02:29:20Z")

</div>

One way of solving this kind of problem is as follows. Instead of preallocating stuff manually, we’ll first create a function `make_buffer` that knows how to create an appropriate work buffer for an in-place function given the function and its input argument. For `f!`, that would be:

```julia
make_buffer(::typeof(f!), p) = similar(p, length(x0))

```

This will use the element type of `p` to create the right type of work buffer. We can then use this to create a function that turns an in-place function into an out-of-place function:

```julia
function make_out_of_place_func(f!, x)
    y = make_buffer(f!, x)
    x -> f!(y, x)
end

```

An out-of-place function created using `make_out_of_place_func` only accepts one element type for `x`, but we want a solution that accepts generic `x`s, so that it works for `ForwardDiff.Dual` as well. One way to do that is as follows:

```julia
struct OutOfPlace{F}
    f!::F
    out_of_place_funcs::Dict{Type, Function}
end
OutOfPlace(f!) = OutOfPlace(f!, Dict{Type, Function}())

eval_f(f::F, x) where {F} = f(x) # function barrier
function (oop::OutOfPlace{F})(x) where {F}
    T = eltype(x)
    f = get!(() -> make_out_of_place_func(oop.f!, x), oop.out_of_place_funcs, T)
    eval_f(f, x)
end

const f = OutOfPlace(f!)

```

i.e., store the in-place function in a `struct`, along with a dictionary that maps element types of `x` to out-of-place functions, where a new out-of–place function is automatically generated if the element type of `x` has not been seen before, but the out-of-place function is reused if a suitable one has already been created. This is pretty efficient:

```julia
using BenchmarkTools, Random
@btime f!(out, p) setup = begin # 199.650 μs (0 allocations: 0 bytes)
    out = similar($x0)
    p = rand!(similar(p0))
end
@btime f(p) setup = p = rand!(similar(p0)); # 220.685 μs (1 allocation: 32 bytes)

```

We can now create an in-place Jacobian function based on the _out-of-place_ version of function `f`:

```julia
# derivative just uses `f`, not `f!`
g! = (outjac, p) -> ForwardDiff.jacobian!(outjac, f, p)

@btime g!(outjac, p) setup = begin
    outjac = Array{Float64}(undef, length(x0), length(p0))
    p = rand!(similar(p0))
end # 388.684 μs (3 allocations: 208 bytes) (allocations would scale with size of `p`)

```

and then use the same trick again:

```julia
# same trick to turn g! into an out-of-place function
make_buffer(::typeof(g!), p) = similar(p, (length(x0), length(p)))
const g = OutOfPlace(g!)

h! = (outjac2, p) -> ForwardDiff.jacobian!(outjac2, g, p)

@btime h!(outjac2, p) setup = begin # 1.637 ms (6 allocations: 576 bytes)
    outjac2 = Array{Float64}(undef, length(x0) * length(p0), length(p0))
    p = rand!(similar(p0))
end

```

I’ve done something similar before in [RigidBodyDynamics.jl/caches.jl at master · JuliaRobotics/RigidBodyDynamics.jl · GitHub](https://github.com/JuliaRobotics/RigidBodyDynamics.jl/blob/master/src/caches.jl#L1)

> **Full code:**
>
> ```julia
> using ForwardDiff
> using LinearAlgebra
> 
> const x0 = randn(50_000);
> const p0 = randn(2);
> f!(out, p) = @. out = p[1] * exp(-x0 * p[2])
> 
> make_buffer(::typeof(f!), p) = similar(p, length(x0))
> 
> function make_out_of_place_func(f!, x)
> y = make_buffer(f!, x) # make_buffer overloads will be defined below
> x -> f!(y, x)
> end
> 
> struct OutOfPlace{F}
> f!::F
> out_of_place_funcs::Dict{Type, Function}
> end
> OutOfPlace(f!) = OutOfPlace(f!, Dict{Type, Function}())
> 
> eval_f(f::F, x) where {F} = f(x) # function barrier
> function (oop::OutOfPlace{F})(x) where {F}
> T = eltype(x)
> f = get!(() -> make_out_of_place_func(oop.f!, x), oop.out_of_place_funcs, T)
> eval_f(f, x)
> end
> 
> const f = OutOfPlace(f!)
> 
> using BenchmarkTools, Random
> @btime f!(out, p) setup = begin # 199.650 μs (0 allocations: 0 bytes)
> out = similar($x0)
> p = rand!(similar(p0))
> end
> @btime f(p) setup = p = rand!(similar(p0)); # 220.685 μs (1 allocation: 32 bytes)
> 
> # derivative just uses `f`, not `f!`
> g! = (outjac, p) -> ForwardDiff.jacobian!(outjac, f, p)
> 
> @btime g!(outjac, p) setup = begin
> outjac = Array{Float64}(undef, length(x0), length(p0))
> p = rand!(similar(p0))
> end # 388.684 μs (3 allocations: 208 bytes) (allocations would scale with size of `p`)
> 
> # same trick to turn g! into an out-of-place function
> make_buffer(::typeof(g!), p) = similar(p, (length(x0), length(p)))
> const g = OutOfPlace(g!)
> 
> h! = (outjac2, p) -> ForwardDiff.jacobian!(outjac2, g, p)
> 
> @btime h!(outjac2, p) setup = begin # 1.637 ms (6 allocations: 576 bytes)
> outjac2 = Array{Float64}(undef, length(x0) * length(p0), length(p0))
> p = rand!(similar(p0))
> end
> 
> ```

---

_[View the full topic](https://discourse.julialang.org/t/nested-forwarddiff-jacobian-calls-with-inplace-function/21232)._