differentiate with respect to x for mutating scalar-valued functions of the form f!(mutating_state, read_only_x, read_only_params)

[RoyCCWang]

I work with Julia functions of the form: f!(state, input, params), where input is the variable that I want to differentiate.

• f! returns a scalar value of type Real.
• The arguments state and params are some custom data type implemented as a struct, which might contain heap-allocated types like Memory or something from AbstractVector.
• The argument input can be of a type from AbstractVector{<:Real}, like `Memory{Float64}.
• The state is mutated when f! is evaluated.
• The contents of params and input are unchanged when f! is evaluated.
• The Julia function f! is type stable, and does not allocate memory when it is evaluated.

I just discovered Mooncake.jl, and got it working via the quick-start section of the readme. I used DifferentiationInterface.jl just to test correctness without worrying allocation and type stability, so I'm using a closure to define f. The rough code looks like

# state, params, x are already defined in scope.

f = xx->query_my_model!(state, xx, params) # since state is mutated, I suspect time-age/"boxing" issues with Julia probably cause this to be non-performant.

using DifferentiationInterface
import Mooncake
backend = AutoMooncake(; config=nothing)
prep = prepare_gradient(f, backend,  x) # x is some input to the model
g = gradient(f, prep, backend,  x)

This works, and my micro benchmark speed shows promise!

I am now going through Mooncake.jl's documentation to see if I can directly differentiate f! without wrapping it in a callable struct or closure, in a type stable and non-allocating manner. I suspect I need to define custom rules for such f! and any mutating functions that is called when f! is evaluated. I'm about to start with the section "Mooncake.jl's Rules System". Does anyone have tips or suggestions on which sections of the documentation to focus on for my usage case?

[willtebbutt]

Hi @RoyCCWang . Thanks for opening this discussion (, and for your issue regarding the docs!) This sounds like an example that Mooncake.jl ought (in principle) to handle fine without any additional rules.

I am now going through Mooncake.jl's documentation to see if I can directly differentiate f! without wrapping it in a callable struct or closure, in a type stable and non-allocating manner.

This sounds like it should be fine. Could you provide an example of what you'd like to do?

I suspect I need to define custom rules for such f! and any mutating functions that is called when f! is evaluated.

You shouldn't usually have to do this.

It would be really very helpful if you could provide a minimal working example (MWE) that demonstrates the kind of thing that you're trying to do / is characteristic of the code you are trying to differentiate. I'm very happy to help with any problems you have, but it's easiest to help if we have some specific code to anchor conversation.

[RoyCCWang]

Hi Will, thanks for the quick reply. My original problem has several nested structs, spanning across multiple custom packages. I'll trace through the data structure of my model, and think about a meaningful MWE. I will work through the entire "Introduction" section of the documentation, then post a MWE here.

[RoyCCWang]

Here is a long example. This code is intended to eventually be statically compiled when trim + smaller binaries arrive, so a few constrained type annotations are used. The goal is to get the derivatives of compute_cost! with respect to the second input slot.

There is a LibModel-prefix for a library model, and a ExtModel-prefix for an user-defined model. The motivation is that there is a few libraries inside a package I'm developing, and there are models or callbacks that a user of the package will want to define outside of the library. Where I commented # arbitrary logic. means some heavy computation occurs there that mutates the state variable for that function scope. In the example I made up some operations that mutate arrays in state. The cost function compute_cost! uses both models. The cost function and both models use the state (potentially mutating), input (where applicable, read-only), and parameters (read-only) setup I mentioned.

The FitVariables struct is for convenience, something like a simplified version of the composite types from ComponentArrays.jl. I create subarrays on-the-fly (i.e., view_b0 and view_noise) to lower the chance of aliasing in downstream algorithms. By the way, I can't find a detailed explanation of how aliasing is handled in the documentation, even though such a detailed discussion on it keeps getting deferred whenever the topic comes up.

using LinearAlgebra, Random

### library area.

struct FitVariables{T <: AbstractFloat}
    contents::Memory{T}

    b_st::Int
    b_fin::Int

    noise_st::Int
    noise_fin::Int

    function FitVariables(::Type{T}, N::Integer) where T <: AbstractFloat
        N > 0 || error("The number of samples must be a positive integer.")

        return new{T}(Memory{T}(undef, N+N), 1, N, N+1, 2*N)
    end
end

function get_num_entries(A::FitVariables)
    return length(A.contents)
end

function view_b0(A::FitVariables)
    return view(A.contents, A.b_st:A.b_fin)
end

function view_noise(A::FitVariables)
    return view(A.contents, A.noise_st:A.noise_fin)
end

function update_vars!(V::FitVariables, x::AbstractVector)
    copy!(V.contents, x)
    return nothing
end

struct LibModelState{T}
    c::Memory{T}

    # buffers
    B::Matrix{T}
    b::Memory{T}

    function LibModelState(::Type{T}, N::Integer) where T <: AbstractFloat
        N > 0 || error("The number of samples must be a positive integer.")

        return new{T}(Memory{T}(undef, N), zeros(T, N, N), Memory{T}(undef, N))
    end 
end

function get_coeffs(A::LibModelState)
    return A.c
end

struct LibModelParams{T}
    A::Matrix{T}
    b0::Memory{T}
    σ::Memory{T}

    function LibModelParams(A::Matrix{T}) where T <: AbstractFloat
        isposdef(A) || error("A is not positive definite.")
        
        return new{T}( copy(A), Memory{T}(undef, size(A,1)), Memory{T}(undef, size(A,1)))
    end 
end

struct CostState{T, PT <: LibModelParams, FT}
    vars::FitVariables{T}

    lib_model_state::LibModelState{T}
    lib_model_params::PT
    
    f_state::FT

    function CostState(A::Matrix{T}, f_state::FT) where {T <: AbstractFloat, FT}
        N = size(A, 1)
        v = FitVariables(T, N)
        s = LibModelState(T, N)
        ps = LibModelParams(A)

        return new{T, typeof(ps), FT}(v, s, ps, f_state)
    end
end

function get_num_vars(A::CostState)
    return get_num_entries(A.vars)
end

# `state` mutates.
function update_state!(state::CostState, x::AbstractVector)

    vars = state.vars
    update_vars!(vars, x)

    copy!(state.lib_model_params.b0, view_b0(vars))
    copy!(state.lib_model_params.σ, view_noise(vars))

    update_model!(state.lib_model_state, state.lib_model_params)

    return nothing
end

# `s`` mutates
function update_model!(s::LibModelState, ps::LibModelParams)

    A, b0, σ = ps.A, ps.b0, ps.σ
    c, B, b = s.c, s.B, s.b

    # arbitrary logic.
    copy!(B, A)
    for (i, j) in Iterators.zip(axes(B,1), eachindex(σ))
        B[i,i] += σ[j]^2
    end
    L = cholesky!(B)

    b .= b0 .* cos.(b)
    ldiv!(c, L, b)
    # cholesky! and ldiv! don't allocate, but can Mooncake (with or without custom rules) avoid allocation for them?

    return nothing
end


struct CostParams{T <: AbstractFloat, MT <: Union{AbstractRange, AbstractVector}, PT, FT}
    obs::Memory{T}
    times::MT

    f_params::PT
    f!::FT
end

function compute_cost!(
    s::CostState, # mutates
    x::AbstractVector{T},
    ps::CostParams,
    ) where T <: AbstractFloat

    f_state = s.f_state
    f!, f_params = ps.f!, ps.f_params

    update_state!(s, x)
    c = get_coeffs(s.lib_model_state)
    
    cost = zero(T)

    ts, ys = ps.times, ps.obs
    for i in eachindex(ts, ys) 
        
        h_t = query_lib_model(c, ts[i])
        fh_t = f!(f_state, h_t, f_params)
        
        cost += (ys[i] - fh_t)^2
    end
    return cost
end

function query_lib_model(c::AbstractVector, t::AbstractFloat)
    return sum( exp(-ci*t^2) for ci in c )
end

### user area.

rng = Random.Xoshiro(0)
T = Float64

struct ExtModelParams{T <: AbstractFloat}
    ws::Memory{T}
    k::T
end

# for intermediate variables that need heap storage.
struct ExtModelState{T <: AbstractFloat}
    a_t::Memory{T}
    A::Matrix{T}

    function ExtModelState(::Type{T}, K::Integer) where T <: AbstractFloat
        a = Memory{T}(undef, K)
        fill!(a, T(NaN))
        return new{T}(a, zeros(T, K, K))
    end
end

struct ExtModelCallable end
function (::ExtModelCallable)(
    s::ExtModelState, # mutates
    t::AbstractFloat,
    ps::ExtModelParams,
    )

    ws, k = ps.ws, ps.k
    a_t, A = s.a_t, s.A

    # arbitrary logic.
    for i in eachindex(ws, a_t)
        a_t[i] = cos(ws[i]*t)
    end

    mul!(A, a_t, a_t')
    return sum( tanh(k*a) for a in A )
end

# make up a template callable.
K = 4
f_params = ExtModelParams(Memory{T}(randn(rng, T, K)), T(1.23))
f_state = ExtModelState(T, K)
f! = ExtModelCallable()

# make up specifications for the RKHS-based warping model.
N = 10
t_st = T(-10.23)
t_fin = T(14.56)
A = randn(rng, T, N, N)
A = A'*A
state = CostState(A, f_state)

N_obs = 30
data_times = LinRange(T(-1), T(5), N_obs)
data_obs = Memory{T}(sinc.(data_times))
params = CostParams(data_obs, data_times, f_params, f!)

# make up a test variable.
N_optim_vars = get_num_vars(state)
x0 = Memory{T}(undef, N_optim_vars)
rand!(rng, x0)

# scale the test input to near the origin, so the derivative magnitude isn't too large.
# If it is much larger than 1, then the numerical and algorithmic differentiation results will be quite different.
x0 .= x0 .* T(1e-3)

# The cost function.
cost_x0 = compute_cost!(state, x0, params)

using FiniteDiff
f = xx->compute_cost!(state, xx, params)

# Numerical derivatives
df! = (aa,xx)->FiniteDiff.finite_difference_gradient!(aa, f, xx)

df_x0_ND = similar(x0)
f_x0 = f(x0)
df!(df_x0_ND, x0)

# Mooncake.jl
using DifferentiationInterface
import Mooncake

backend = AutoMooncake(; config=nothing)
prep = prepare_gradient(f, backend, x0)
df_x0_MC = gradient(f, prep, backend, x0)

# relative discrepancy between the numerical and Mooncake's derivative.
println("Check correctness:")
@show norm(df_x0_ND-df_x0_MC)/norm(df_x0_ND)
println("column 1: numerical, column 2: MoonCake")
display([df_x0_ND df_x0_MC])
println()

On my machine, I have the following in my REPL:

Check correctness:
norm(df_x0_ND - df_x0_MC) / norm(df_x0_ND) = 2.5088724298094027e-7
column 1: numerical, column 2: MoonCake
20×2 Matrix{Float64}:
  1621.97          1621.97
 -2408.44         -2408.44
  -546.831         -546.83
   793.561          793.561
  -170.088         -170.088
  -719.753         -719.753
    86.0113          86.0116
  -497.114         -497.114
 -3956.18         -3956.19
  -844.135         -844.135
     0.00590783       0.0059273
     0.0146209        0.0146209
     0.000206735      0.000206736
     0.00295211       0.00295211
     2.43127e-6       2.43232e-6
     0.000147365      0.000147367
     3.31278e-5       3.31269e-5
     0.000208711      0.00020871
     0.0647075        0.0647075
     0.0008313        0.000831298

  17.543 μs (0 allocations: 0 bytes)
  17.603 μs (1 allocation: 16 bytes)
  702.296 μs (104 allocations: 1.80 KiB)
  4.375 ms (3396 allocations: 141.43 KiB)

In this case, the non-in-place gradient seems to produce a result df_x0_MC that is close to the numerical derivative df_x0_ND. Now, benchmark compute_cost! and the derivatives:

using BenchmarkTools
@btime compute_cost!($state, $x0, $params)
@btime $f($x0)
@btime $df!($df_x0_ND, $x0)
@btime gradient($f, $prep, $backend, $x0)

On my machine, my REPL shows:

  17.543 μs (0 allocations: 0 bytes)
  17.603 μs (1 allocation: 16 bytes)
  702.296 μs (104 allocations: 1.80 KiB)
  4.375 ms (3396 allocations: 141.43 KiB)

If I don't use

using DifferentiationInterface
import Mooncake

backend = AutoMooncake(; config=nothing)
prep = prepare_gradient(f, backend, x0)
df_x0_MC = gradient(f, prep, backend, x0)

and use only Mooncake.jl commands, how can I directly differentiate compute_cost!,

Currently, I need to use a closure f! or some callable struct wrapper around compute_cost!, the mutable buffer state, and read-only params. If I only use Mooncake.jl commands, willing to write some rules and use a pre-allocatedstate_shadow = deepcopy(state) to shadow state and params, how can I directly differentiate compute_cost! such that the resultant reverse AD function is type stable and doesn't allocate?

[willtebbutt]

Thanks for this example. In short: you really ought not to have to write any rules to make the above code fast -- it's something that should "just work" and be performant.

Plainly this isn't currently the case. I'm going to dig into it a bit -- it looks like something is going quite wrong in the reverse pass, so I need to figure out understand what's going on.

[willtebbutt]

Okay. Version 0.4.53 gives the following benchmarking results on my machine:

julia> @btime compute_cost!($state, $x0, $params)
  10.166 μs (0 allocations: 0 bytes)

julia> @btime $f($x0)
  10.166 μs (1 allocation: 16 bytes)

julia> @btime $df!($df_x0_ND, $x0)
  410.750 μs (104 allocations: 1.80 KiB)

julia> @btime gradient($f, $prep, $backend, $x0)
  67.209 μs (2103 allocations: 99.05 KiB)

This version fixed a couple of performance bugs that were really hurting performance in your case.

The remainder of the allocations are associated to things which we can't yet get rid of. I'll explain further in an issue.
that I'll open shortly.

[willtebbutt]

Issue in question is #403 .

[RoyCCWang]

Amazing, thanks! I'll post comments on the auto-generation of a stateful rrule!! in the issue you referenced. I'll leave this discussion unanswered for now.