Skip to content

Avoid batch size of 0 for empty inputs #835

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 3 commits into
base: main
Choose a base branch
from

Conversation

penelopeysm
Copy link
Contributor

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.

@penelopeysm penelopeysm requested a review from gdalle as a code owner August 15, 2025 00:16
@gdalle
Copy link
Member

gdalle commented Aug 16, 2025

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
Copy link
Member

gdalle commented Aug 16, 2025

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

@gdalle gdalle marked this pull request as draft August 16, 2025 09:27
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants