Add batching pass - Githubissues

EnzymeAD / Enzyme

High-performance automatic differentiation of LLVM and MLIR.

https://enzyme.mit.edu

Other

1.2k stars 100 forks source link

Closed wsmoses closed 1 week ago

wsmoses commented 1 week ago

cc @mofeing @avik-pal

mofeing commented 1 week ago

I'm not sure how batching works in Enzyme, but some StableHLO ops have support for batching (e.g. dot_general). Should this pass modify those ops?

wsmoses commented 1 week ago

Eventually, yes but for now this just needs to operate on scalars (to implement broadcasting properly in reactant)