search

MasonProtter / Bumper.jl

Bring Your Own Stack
MIT License
152 stars 6 forks source link

Add some EnzymeRules #21

Open MasonProtter opened 9 months ago

MasonProtter commented 9 months ago

Currently, Enzyme.jl's reverse mode autodiff doesn't work correctly with Bumper.jl because if you give it a Duplicated buffer, it'll += accumulate results into the duplicated buffer making the answer depend on the state of the buffer at the start of the program.

It'd be good if we could set up some EnzymeRules to explicitly teach Enzyme how to handle Bumper.jl allocations and deallocations. I don't really know how to do this though, so if anyone wants to take it on, or work on it together please do.

MilesCranmer commented 6 months ago

Just so I understand, you could still use bump allocation internal to a differentiated function, right? You are basically wondering if it could also be done with bump allocation wrapping an enzyme call?

MasonProtter commented 6 months ago

I'm pretty sure it'll not work even if you do it internally (unless you manually zero out your buffer)