Enzyme wrappers for AD utilities in Lux

[avik-pal]

Functions

• [ ] batched_jacobian
  • [ ] Enzyme.Forward
  • [ ] Enzyme.Reverse
• [ ] vector_jacobian_product
• [ ] jacobian_vector_product

Nested Autodiff

Enzyme needs us to use autodiff_thunk for the internal call, so we can easily replace the call inside custom rrules.

• [ ] EnzymeRules
• [ ] CRC.rrule