Specific preparation values for `DifferentiationInterfaceTest.Scenario`s

[penelopeysm]

Hello @gdalle!

I was wondering if there is any way to specify the values at which the gradient preparation is carried out when constructing a Scenario. For example, this would be useful to (reproducibly) demonstrate that compiled ReverseDiff can potentially give incorrect values:

import DifferentiationInterface as DI
using ReverseDiff, ADTypes

f(x) = x[1] > 0 ? x[2] : x[2]^2

prep = DI.prepare_gradient(f, AutoReverseDiff(; compile=true), [5.0, -5.0])

DI.value_and_gradient(f, prep, AutoReverseDiff(; compile=true), [-5.0, 5.0])
# actual: (5.0, [0.0, 1.0]) - should of course be: (25.0, [0.0, 10.0])

[gdalle]

Hi Penelope,

I didn't know anyone else was actually using DifferentiationInterfaceTest, that's why the interface is still a bit raw.
Your suggestion is a good one, and it would also solve cases where we need a "variant" of the scenario. One such case is adapting preparation results when they were computed for smaller or larger inputs.

Perhaps the Scenario type could have a field named prepare_on::Scenario where we sum up the values needed at preparation time? Alternately, we could add 4 fields x_prep, y_prep, t_prep and contexts_prep. What do you think?

[penelopeysm]

I didn't know anyone else was actually using DifferentiationInterfaceTest

We are not actually using it 🥲 (the tl;dr is that we want to test LogDensityProblems.logdensity_and_gradient for Turing models, but you might remember from TuringLang/DynamicPPL.jl#806 that there are several bits of code between that and DI.value_and_gradient, so using DITest would require duplicating all those steps when constructing scenarios.)
However, the need for that came up when writing those tests and I figured that I should put in an issue too haha.

You're the expert here, but as a potential user I might be inclined towards extra fields x_prep, ... - that seems the most straightforward / customisable option (although it'd be ugly having to check isnothing on all of them)

[gdalle]

With #786 you can do:

using DifferentiationInterface, DifferentiationInterfaceTest
using ForwardDiff: ForwardDiff
using ReverseDiff: ReverseDiff

f(x) = x[1] > 0 ? x[2] : x[2]^2

scen = Scenario{:gradient,:out}(
    f, [-5.0, 5.0]; prep_args=(; x=[5.0, -5.0], contexts=()), res1=[0.0, 10.0]
)

backends = [
    AutoForwardDiff(), AutoReverseDiff(; compile=false), AutoReverseDiff(; compile=true)
]

test_differentiation(backends, [scen])

and get the expected results:

Test Summary:                                                 | Pass  Fail  Total  Time
Testing correctness                                           |   60    12     72  0.2s
  AutoForwardDiff()                                           |   24           24  0.0s
  AutoReverseDiff()                                           |   24           24  0.0s
  AutoReverseDiff(compile=true)                               |   12    12     24  0.2s
    gradient                                                  |   12    12     24  0.2s
      Scenario{:gradient,:out} f : Vector{Float64} -> Float64 |   12    12     24  0.2s
        Correctness                                           |   12    12     24  0.2s