Used New Fused Ops from LuxLib

LuxDL / Lux.jl

Elegant & Performant Scientific Machine Learning in Julia

https://lux.csail.mit.edu/

MIT License

475 stars 57 forks source link

Closed avik-pal closed 4 months ago

avik-pal commented 4 months ago

[x] Fix Nested AD on CUDA
[x] ~ReverseDiff and Tracker seems to be slowed down. Rewrite the rrule_via_ad to directly use ForwardDiff~ they don't hit the fused kernels anyways