First test scenarios for Flux gradients

[nialamarcotte]

Work done during JuliaCon Hackathon; preparing test for neural networks

[gdalle]

Thanks so much for your contribution! It may take a few days but I'll get back to you, give you some feedback on the code and take it from here if that's okay. I'm sorry that the task was so daunting, but very grateful that you agreed to give it a try :) And if you need any more resources on AD, let me know

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 98.23009% with 2 lines in your changes missing coverage. Please review.

Project coverage is 96.58%. Comparing base (63fa856) to head (ceeb25f).
Report is 2 commits behind head on main.

Files	Patch %	Lines
DifferentiationInterfaceTest/src/utils/misc.jl	71.42%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #352      +/-   ##
==========================================
+ Coverage   96.57%   96.58%   +0.01%     
==========================================
  Files          98       99       +1     
  Lines        4813     4888      +75     
==========================================
+ Hits         4648     4721      +73     
- Misses        165      167       +2     

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

[gdalle]

CI

• Adjust groups

DIT source

• Support custom isequal function in correctness tests, in addition to isapprox
• Adjust mycopy_random and mysimilar for arbitrary structures with Functors

DIT extensions

• Add Flux extension with gradient scenarios, whose correctness is checked against FiniteDifferences (see DifferentiationInterface testing FluxML/Flux.jl#2469 (comment))
• Add more scenarios with various types of networks

DIT tests

• Test new Flux scenarios with Zygote

DI extensions

• In the Enzyme extension, support non-array inputs in reverse mode:
  • Relax type constraints
  • Avoid querying eltype(x), backpropagate dy=true instead

DI tests

• Test new Flux scenarios with Zygote (Enzyme fails)

[gdalle]

@CarloLucibello I got the first batch of tests working with Zygote. The syntax is a bit counter-intuitive because DI doesn't yet support multiple arguments, so I have to put the data x inside the loss closure, but it gets the job done for now.
The Enzyme tests error on two cases though, the ones involving Flux.Scale and Flux.SkipConnection. Any idea what could be different about those two? @wsmoses could this be one of the cases where we want to Duplicate the function f itself for the time being?
What other backends do we want to test here? Tracker?

[wsmoses]

Actually this is probably a case of exactly the opposite in which doing so would be harmful. Marking the closure as differentiable would mean that Enzyme would also end up computing the derivative of the input [and then throwing it away].

[wsmoses]

Per the error message, it looks like the closure is causing issues where it can't be compiled due to mixed activity [e.g. one element of the closure being differentiable the other not]. As the error message states, you can enable runtime activity to resolve at a slight [but not asymptotic performance loss]. But also using multiple arguments here in Enzyme directly would also presumably resolve this

[gdalle]

Per the error message, it looks like the closure is causing issues where it can't be compiled due to mixed activity [e.g. one element of the closure being differentiable the other not]. As the error message states, you can enable runtime activity to resolve at a slight [but not asymptotic performance loss].

I tried enabling runtime activity but the stack traces just got angrier 😱

But also using multiple arguments here in Enzyme directly would also presumably resolve this

Agreed, but setting this up in the single-argument case also helped me resolve some other issues to pave the way for non-array testing, so it's a great first step. For the multi-argument case, what I need most of all is inspiration and ideas to improve my initial API proposal at #311 (comment). The idea is to be slightly different than Enzyme, and closer to the mathematical concepts that users may be familiar with.

[gdalle]

Merging this without the Enzyme tests because there's already some useful stuff in there. Billy if you have a clue what the problem was in https://github.com/gdalle/DifferentiationInterface.jl/actions/runs/9963932557/job/27531092106 with runtime activity switched on, I'd love to hear your thoughts