feat: access overloaded inputs from preparation result

[ErikQQY]

Fix: #668

This PR only adds the overloaded_inputs functionalities for Jacobian preparations, not sure if the derivative, gradient and hessian preparation need this API.

The tests for this feature maybe too simple, need some suggestions for better test cases.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 97.61%. Comparing base (240e7e8) to head (805b000).
Report is 1 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #672   +/-   ##
=======================================
  Coverage   97.61%   97.61%           
=======================================
  Files         112      114    +2     
  Lines        5610     5623   +13     
=======================================
+ Hits         5476     5489   +13     
  Misses        134      134           

Flag	Coverage Δ
DI	98.76% <100.00%> (+<0.01%)	⬆️
DIT	95.35% <ø> (ø)

Flags with carried forward coverage won't be shown. Click here to find out more.

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

[ErikQQY]

CI failures seem not related.

[gdalle]

Thanks for getting this started! In terms of design, I think we should not just return an eltype but instead we should return the full dualized array. The reason being that some backends (like Mooncake) don't take an elementwise dual.
Could you also make sure that code coverage does not decrease?

[ErikQQY]

All tests passed now.

[ErikQQY]

Also implemented the overloaded_inputs for derivative and gradient preparations.

[ErikQQY]

I believe this PR is ready now

[gdalle]

thanks, I'll take a look!

[gdalle]

@ErikQQY I have given this some more thought and I'm not sure the abstraction we picked is the correct one.

Symbolic backends also rely on operator overloading, but they will only call the function on "overloaded inputs" once, during preparation. This serves to compile a fast version of the derivative, which is then applied to Float64 inputs. So it is very different from ForwardDiff and ReverseDiff, where the "overloaded inputs" are used every time.

In general, we have no guarantee that the function f is called, either during preparation, or during prepared or unprepared differentiation.

[gdalle]

I tried to express this in the docstring but it may not be precise enough

[ErikQQY]

Oh we are returning the eltype of the dualized array now 👍.

Symbolic backends also rely on operator overloading, but they will only call the function on "overloaded inputs" once, during preparation.

So we also need the overload_input_type for symbolic backends preparation?

[gdalle]

Not exactly, those actually are two different concepts.

• For ForwardDiff and ReverseDiff, the overloaded type is used at every call to the differentiation routine, which is obtained by going through the function itself.
• For Symbolics, the overloaded type is only used once during preparation, after which the initial function is discarded and runtime-generated functions are used instead which directly encode the derivatives, hence they can be applied to Float64 directly.
• For source transformation backends like Zygote and Enzyme, the function may or may not be called, but there are no overloaded types to speak of.

I guess I'm wondering how to get this point across through a better choice of terminology. And I'm also wondering whether we can find yet another solution for this problem, because introducing something which behaves so differently across backends sounds like a footgun.

[gdalle]

Let's merge this, see how it goes downstream, and then solidify the interface if necessary. It's explicitly marked as not public for the time being so there's no harm