fix: improve support for empty inputs (still not guaranteed)

[penelopeysm]

I ran into #802 again while refactoring some tests on Bijectors.jl, so this is a patch that lets forward-mode Enzyme work with empty inputs.

[gdalle]

I gave it a try too and it's a bit harder than it looks to make this work consistently. First, there are unavoidable discrepancies across backends:

• ForwardDiff can deal with zero chunk sizes, in fact it chooses them voluntarily with ForwardDiff.pickchunksize(0).
• Enzyme cannot deal with zero chunk sizes, it will error here.
• Other backends must have a default chunk size of 1, we cannot use max(1, length(x)) or it becomes type-unstable (see here).

This suggests we need to handle batch sizes of 0 while also setting a minimum of 1 when we can.
Second, there are some difficulties related to the DI operators themselves:

• For wrong-mode pushforwards and pullbacks, we propagate basis vectors, which error when empty (e.g. here). This is solvable.
• For Jacobians and Hessians, we precompute a set of basis vectors and perform preparation on the first one, which may not exist when the vector is empty (e.g. here). This is solvable.
• For Jacobians and Hessians, mapreduce(hcat, collection) errors when the collection is empty (e.g. here). I'm not sure how to solve that one.

[gdalle]

To give you an example of the kind of changes necessary, you can take a look at this commit. I'm not sure supporting empty arrays is worth it, especially if support is always gonna be divergent across backends as you noted in #802

[penelopeysm]

Thanks for looking into it in more detail! I definitely see the issues with making it consistent. I didn't really set out to solve it for all backends, that was a task that was too annoying and possibly also too unimportant -- hence my very tiny single commit 😄. I don't think that DI needs to explicitly guarantee or document that empty inputs will work for all backends; but it's just nice to not get an error and if some code could be tweaked to make that happen, that's a net benefit IMO.

Did that commit alone break ForwardDiff or some other backend? I was sort of hoping that CI would catch issues that I hadn't thought of or didn't know of (of which I'm sure there are plenty). But also my impression was that the lines I changed should only affect empty inputs.

[codecov]

Codecov Report

❌ Patch coverage is 96.66667% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 88.12%. Comparing base (4a59a3a) to head (22bd82e).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
DifferentiationInterfaceTest/test/zero_backends.jl	66.66%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #835      +/-   ##
==========================================
+ Coverage   78.24%   88.12%   +9.87%     
==========================================
  Files         128      132       +4     
  Lines        7805     7872      +67     
==========================================
+ Hits         6107     6937     +830     
+ Misses       1698      935     -763     

Flag	Coverage Δ
DI	85.21% <100.00%> (+13.70%)	⬆️
DIT	95.69% <90.00%> (-0.03%)	⬇️

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.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[gdalle]

Did that commit alone break ForwardDiff or some other backend?

I don't think so but I wanted to add proper tests covering this line, and to make sure that the global approach is at least vaguely coherent. I tried defining a batch size of 0 for empty arrays by default, but I think you're right that 1 is better. I ended up adding a bunch of empty test scenarios in DIT to make my life easier. DI's default implementation of jacobian still fails on empty vectors because of the uninitialized mapreduce, but apart from that this is more satisfactory to me. Can you check whether the added tests cover everything that was failing for you?

[gdalle]

Enzyme tests failing due to EnzymeAD/Enzyme.jl#2523

DIT tests failing because I must have introduced a type instability somewhere