Specific preparation values for DifferentiationInterfaceTest.Scenarios

[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)