Add `vector_jacobian_product` and `jacobian_vector_product` functions - Githubissues

LuxDL / Lux.jl

Elegant & Performant Scientific Machine Learning in Julia

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

MIT License

456 stars 54 forks source link

Add `vector_jacobian_product` and `jacobian_vector_product` functions #623

Closed avik-pal closed 2 months ago

avik-pal commented 2 months ago

[x] Fixes #621
[x] Add rrules
- [x] VJP
- [x] JVP
[x] Add tests
[x] Add a manual entry: Maybe showcasing Hutchinson Trace Estimator