Improvements to Batched Jacobian

LuxDL / Lux.jl

Elegant & Performant Scientific Machine Learning in Julia

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

MIT License

456 stars 54 forks source link

Closed avik-pal closed 2 months ago

avik-pal commented 2 months ago

[x] Reverse Batched Jacobian
- [x] Forward Pass
- [x] Nested Part
[x] Use Tasks for ForwardDiff batched_jacobian
- [x] Check if it fails for cudnn