Single-pass value, derivative and second derivative

[gerlero]

Hi @gdalle. After discovering this new package, I'm considering a possible migration from AbstractDifferentiation in my projects... I wonder, is there a way to get the value and first two derivatives in a single pass for an $\mathbb{R}\to\mathbb{R}$ function (i.e., an equivalent to AbstractDifferentiation's value_derivative_and_second_derivative)?

[gdalle]

Hi @gerlero,
I remember your PR to AbstractDifferentiation implementing this functionality (JuliaDiff/AbstractDifferentiation.jl#122). At the moment, DifferentiationInterface doesn't offer it, and second-order operators have to be computed separately.
I've been meaning to add value_gradient_and_hessian for a while, but only two backends have optimized implementations so far: ForwardDiff and ReverseDiff. As for value_derivative_and_second_derivative, only ForwardDiff can make it in a single pass (that I'm aware of). It's definitely possible, but I haven't gotten around to it, precisely because you get no benefits at all except for a single backend. I think TaylorDiff might also be efficient in this case but it's far from mature.
The other issue is that DifferentiationInterface computes derivative indifferently for $\mathbb{R} \to \mathbb{R}$ and for $\mathbb{R} \to \mathbb{R}^n$. I assume your ForwardDiff implem only works for the former?

[adrhill]

I've been meaning to add value_gradient_and_hessian for a while, but only two backends have optimized implementations so far

I think our baseline for performance should be the two separate calls to value_and_gradient and hessian (or gradient and value_and_hessian) that are currently necessary for users.

[gerlero]

Thanks. I'd be very interested in using it if it were available (even if only supported by a single backend), especially given that SciML is now using this package and I already depend on a lot of SciML.

The other issue is that DifferentiationInterface computes derivative indifferently for and for . I assume your ForwardDiff implem only works for the former?

Yes. I've just checked it and indeed my implementation in AbstractDifferentiation doesn't return a sensible result for non-scalar valued functions (not saying it couldn't be made to work though).

[gdalle]

The main obstacle to adding new operators to DifferentiationInterface is not really the operator itself, it's writing the proper tests and running them every time. I think for this one it shouldn't be too much trouble though.
Here's an example implementation that you can test in your code, I'll try adding it later to see how easily I can get away with it.

using ADTypes
using DifferentiationInterface
using DifferentiationInterface: inner
using ForwardDiff: ForwardDiff

## Code

struct FirstAndSecondDerivativeExtras{E1,E2}
    extras1::E1
    extras2::E2
end

function prepare_derivative_and_second_derivative(f, backend::SecondOrder, x)
    extras1 = prepare_derivative(f, inner(backend), x)
    extras2 = prepare_second_derivative(f, backend, x)
    return FirstAndSecondDerivativeExtras(extras1, extras2)
end

function value_derivative_and_second_derivative(f, backend::SecondOrder, x, extras=prepare_derivative_and_second_derivative(f, backend, x))
    y, der = value_and_derivative(f, inner(backend), x)
    der2 = second_derivative(f, backend, x)
    return y, der, der2
end

## Example

backend = SecondOrder(AutoForwardDiff(), AutoForwardDiff())

f(x) = [one(x), x, x^2, x^3]

extras = prepare_derivative_and_second_derivative(f, backend, 1.0);

y, der, der2 = value_derivative_and_second_derivative(f, backend, 2.0, extras)

[gerlero]

Interesting. For ForwardDiff, can this package also accept a special case so that it uses the same single-pass implementation as AbstractDifferentiation does?

[gdalle]

Actually disregard the code above, there's a faster way to do this which I'm currently writing in the package.
DifferentiationInterface works like AbstractDifferentiation: backends are free to add as many special cases as they want in the extensions. So once I get the basic fallbacks up and running (and tested!), I'll write the special case for ForwardDiff to do it in one pass instead of two.

[gerlero]

As of v0.5.5, DifferentiationInterface can provide values and first and second derivatives in one pass with the AutoForwardDiff() backend via the value_and_derivative and value_derivative_and_second_derivative functions.

Thanks to @gdalle for his willingness and help in adding this functionality to the package.