In-place versions e.g. `logdensity_and_gradient!`?

[penelopeysm]

In DynamicPPL/Turing, we quite often wrap models inside something that obeys the LDP interface. However, under the hood we use DifferentiationInterface.value_and_gradient to calculate logdensity_and_gradient.

It's very trivial to extend this code to support in-place versions, by changing from DI.value_and_gradient to DI.value_and_gradient!. Similar considerations apply for the second-order functions too. However, LDP doesn't provide a logdensity_and_gradient! function, which would be the natural place to stick an overload.

Conscious of the comment in #110 (comment)

I am happy to add API for anything if someone actually needs it. If there is a project, however experimental, that needs this, please comment here or make a PR.

I'd mention that I've seen at least two instances where people have manually implemented what is essentially logdensity_and_gradient!, and this is just off the top of my head:

https://github.com/Julia-Tempering/Pigeons.jl/blob/115265329bb291c389349ab05d9298f17183ba0c/ext/PigeonsDynamicPPLExt/interface.jl#L109-L116

https://github.com/TuringLang/AdvancedVI.jl/blob/721964c80d8d98ee8cb3b12466a257449c7d856e/src/algorithms/fisherminbatchmatch.jl#L94-L97

[tpapp]

Technically, it would be trivial to add an in-place method. There are some decisions to make, such as whether to advertise that in capabilities (and whether to make that a breaking change), and whether to have a fallback (the in-place method to the functional one, or vice versa). But these can be discussed and overcome.

Originally, I did not include an in-place method because for the kinds of problems I am working with, the cost of a single allocation is negligible compared to the cost of the log posterior and AD for any nontrivial problems. Personally I think it is a premature optimization, but I am open to benchmarks.

[penelopeysm]

Right -- if it were me making a PR, I would probably not change capabilities, and add a definition of value_and_gradient! to fall back to value_and_gradient plus copyto!. The rationale being that from an interface perspective it doesn't matter whether the internal implementation is efficient (saving the allocation) or not.

I do see the point about it being a very small gain compared to AD; but also it's a comparatively small change to LogDensityProblems. I benchmarked on eight-schools. I think probably that still falls under what you'd consider the trivial category, so I'm not really trying to use it to make a point -- just wanted to have some numbers so that we have something to talk about

• primal: 175 ns
• gradient with reverse Enzyme: 530 ns
• in-place gradient with reverse Enzyme: 500 ns

Benchmarking code

using DynamicPPL, Distributions, LinearAlgebra, ADTypes, Enzyme, Chairmarks, LogDensityProblems
using DynamicPPL: _get_input_vector_type, _use_closure, LogDensityAt, logdensity_at
import DifferentiationInterface as DI

J = 8
y = [28, 8, -3, 7, -1, 1, 18, 12]
sigma = [15, 10, 16, 11, 9, 11, 10, 18]

@model function esc(J, y, sigma)
    mu ~ Normal(0, 5)
    tau ~ truncated(Cauchy(0, 5); lower=0)
    theta ~ MvNormal(fill(mu, J), tau^2 * I)
    for i in 1:J
        y[i] ~ Normal(theta[i], sigma[i])
    end
end

function logdensity_and_gradient!(
    ldf::LogDensityFunction, out::AbstractVector, params::AbstractVector{<:Real}
)
    # `params` has to be converted to the same vector type that was used for AD preparation,
    # otherwise the preparation will not be valid.
    params = convert(_get_input_vector_type(ldf), params)
    return if _use_closure(ldf.adtype)
        DI.value_and_gradient!(
            LogDensityAt(
                ldf.model,
                ldf._getlogdensity,
                ldf._varname_ranges,
                ldf.transform_strategy,
                ldf._accs,
            ),
            out,
            ldf._adprep,
            ldf.adtype,
            params,
        )
    else
        DI.value_and_gradient!(
            logdensity_at,
            out,
            ldf._adprep,
            ldf.adtype,
            params,
            DI.Constant(ldf.model),
            DI.Constant(ldf._getlogdensity),
            DI.Constant(ldf._varname_ranges),
            DI.Constant(ldf.transform_strategy),
            DI.Constant(ldf._accs),
        )
    end
end

model_esc = esc(J, y, sigma)
ad_en = AutoEnzyme(; mode=Enzyme.set_runtime_activity(Enzyme.Reverse))
ldf = LogDensityFunction(model_esc, getlogjoint_internal, LinkAll(); adtype=ad_en)
x = rand(ldf)
y = similar(x)

@be LogDensityProblems.logdensity($ldf, $x)
@be LogDensityProblems.logdensity_and_gradient($ldf, $x)
@be logdensity_and_gradient!($ldf, $y, $x)

[tpapp]

Thanks for making a self-contained benchmark, this is pretty much in line with my experience.

from an interface perspective it doesn't matter whether the internal implementation is efficient (saving the allocation) or not

You make a valid point, but I would guess that the moment this option is available, some users would submit issues and PRs to make the implementations more efficient. The implementation cost would not be born by this package, as it only defines an API, but packages like LogDensityProblemsAD.jl that do the actual implementation.

Again, I am not against this is someone really needs it, I am just unconvinced about the practical utility. IMO it would also make extensions like #112 more complicated (but that can be worked around too of course).

[penelopeysm]

That is a fair assessment 🙂