Leverage function_annotation for AutoEnzyme

[gdalle]

DI extensions

Enzyme: adapt to SciML/ADTypes.jl#77 with the new function_annotation kwarg:

• Pass get_f_and_df(f, backend) to every operator with the following behavior with the following behavior:
  • if function_annotation <: Nothing, pass f
  • if function_annotation <: Const, pass Const(f)
  • if function_annotation <: Duplicated, pass Duplicated(f, make_zero(f))
  • BatchDuplicated on f not necessary yet but it will be once I implement batched operators
• Restrict the use of high-level API like Enzyme.gradient and Enzyme.jacobian to cases where f is either not annotated or Const: function_annotation <: Union{Nothing,Const}
• Emulate the new default behavior of autodiff when calling autodiff_thunk:
  • force Const(f) if f is not annotated
  • make sure that Enzyme throws an error if something is written inside of f with EnzymeCore.set_err_if_func_written

DI tests

Enzyme:

• Test that function_annotation=Nothing and function_annotation=Const work on normal scenarios
• Test that function_annotation=Duplicated works on closurized scenarios
• In SecondOrder tests, for Enzyme.Reverse over ForwardDiff, mark f as Const because the analysis is too conservative

---

@wsmoses would you add any other annotation? We'll worry about performance later (e.g. zero-ing in place instead of re-creating make_zero(f)).

Related:

• Duplicating closures EnzymeAD/Enzyme.jl#1669
• Make Enzyme dispatches compatible with closures  #339

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 96.70%. Comparing base (0e5ecfd) to head (bf0dee3).

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #407      +/-   ##
==========================================
+ Coverage   96.68%   96.70%   +0.01%     
==========================================
  Files         103      103              
  Lines        4981     5003      +22     
==========================================
+ Hits         4816     4838      +22     
  Misses        165      165              

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

[gdalle]

Some of the failures I see here are for autodiff_thunk, which apparently needs a specific activity and cannot get away with just receiving f. Any suggestions @wsmoses?

[wsmoses]

@gdalle it's probably worth testing that an error is thrown for closure scenaios that write (and nor for ones that don't)

[gdalle]

@wsmoses I think this is ready, care to take a second look?

@gdalle it's probably worth testing that an error is thrown for closure scenaios that write (and nor for ones that don't)

At the moment I can't do that easily because my function test_differentiation runs a test set, so I can't use @test_throws on it. The exceptions would just show up as failed tests within the global test suite.

[gdalle]

I'm gonna merge this because the new version of Enzyme causes errors in my existing tests without it (see e.g. the current state of #329). We can hash out the details later