feat: VJP utility based on autodiff_thunk

[gdalle]

Fixes #1853

Related:

• DI and Enzyme JuliaDiff/DifferentiationInterface.jl#721

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 75.43%. Comparing base (037dfed) to head (6bf6299).
Report is 353 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2309      +/-   ##
==========================================
+ Coverage   67.50%   75.43%   +7.92%     
==========================================
  Files          31       56      +25     
  Lines       12668    16758    +4090     
==========================================
+ Hits         8552    12642    +4090     
  Misses       4116     4116              

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[wsmoses]

Just wondering why not just add a vjp, that seems similar to this except doesn't need separate ones for batches

[gdalle]

Indeed I'm not a fan of the name seeded_autodiff_thunk either. vjp does convey the notion that input and output have to be vectors, so maybe pullback is more generic? What kind of signature do you have in mind?

[wsmoses]

I'm personally fine with the name vjp but at minimum go ahead and add the code and we can iterate on names in parallel

[gdalle]

wait what do you mean by "add the code"? I thought your request of "adding a vjp" was mostly about naming? what is missing here?

[wsmoses]

We shouldn't have two separate functions here for batched vs not

[gdalle]

I also thought of doing it via dispatch but how do we handle the case where dresult itself is supposed to be an NTuple (as opposed to a batch)?

[vchuravy]

In some way, this feels equivalent to implementing autodiff but for non-scalar returns and fixing my old mistake of always passing in one(T) as the seed.

I think of autodiff(Reverse generally as vjp (with the convention that the output is updated in-place)

[gdalle]

We could also call it autodiff but then we'd need to figure out how the output seed is passed. It can't be inside a Duplicated or BatchDuplicated because we have no primal

[vchuravy]

It can't be inside a Duplicated or BatchDuplicated because we have no primal

Seed and BatchSeed? Just throwing some ideas into the air.

[gdalle]

And how would you see the order of arguments?

autodiff(Reverse, f, seed, args...)
autodiff(Reverse, f, args...; seed=...)

[vchuravy]

The first variant, since that is already the convention used for the activity of the return.