Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Batched pushforward, pullback and hvp #320

Merged
merged 23 commits into from
Jun 20, 2024
Merged

Batched pushforward, pullback and hvp #320

merged 23 commits into from
Jun 20, 2024

Conversation

gdalle
Copy link
Member

@gdalle gdalle commented Jun 17, 2024

DI source

  • Define Batch{B,T} to wrap an NTuple, and pick_batchsize(backend, n)
  • Implement batched versions of pushforward, pullback and hvp (the seed dx/dy is a Batch, and the result too)
  • Completely revamp dense Jacobian and Hessian to use batched operators
  • Completely revamp sparse Jacobian and Hessian to use batched operators
  • Pick batch size based on the dense backend for Jacobian, or dense outer backend for Hessian
  • Reduce code duplication in Jacobian between f and f! by splatting f_or_f!y
  • Define internal backends AutoForwardFromPrimitive and AutoReverseFromPrimitive to pretend that a given backend only defines value_and_pushforward / value_and_pullback => tests our entire cascade of fallbacks

DI extensions

  • Implement DI.pick_batchsize with ForwardDiff's internal mechanism
  • Use pick_batchsize in Enzyme extension
  • Use forward-over-reverse for batched HVP in Zygote extension (reverse over reverse made Julia lag forever)

DI tests

  • Test our entire cascade of fallbacks with
    • AutoForwardFromPrimitive + ForwardDiff
    • AutoReverseFromPrimitive + ReverseDiff

TODO (in another PR) => see #312

  • Test Batch struct / pick_batchsize
  • Check type stability and performance
  • Test batched versions of the operators for themselves, and not only as part of jacobian or hessian (will be easier thanks to Revamp test scenarios #323)
  • Implement batched mode in ForwardDiff and Enzyme extensions (+ maybe FiniteDiff?)

@codecov-commenter
Copy link

codecov-commenter commented Jun 17, 2024

Codecov Report

Attention: Patch coverage is 98.57820% with 6 lines in your changes missing coverage. Please review.

Project coverage is 96.51%. Comparing base (004e934) to head (e6c05f1).

Files Patch % Lines
DifferentiationInterface/src/sparse/jacobian.jl 98.23% 2 Missing ⚠️
...dDiffExt/DifferentiationInterfaceForwardDiffExt.jl 75.00% 1 Missing ⚠️
...fferentiationInterface/src/first_order/jacobian.jl 98.96% 1 Missing ⚠️
...fferentiationInterface/src/second_order/hessian.jl 96.55% 1 Missing ⚠️
DifferentiationInterface/src/sparse/hessian.jl 97.43% 1 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main     #320      +/-   ##
==========================================
+ Coverage   96.48%   96.51%   +0.03%     
==========================================
  Files          95       96       +1     
  Lines        4435     4564     +129     
==========================================
+ Hits         4279     4405     +126     
- Misses        156      159       +3     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

@gdalle gdalle marked this pull request as ready for review June 20, 2024 08:47
@gdalle gdalle merged commit f62c9dd into main Jun 20, 2024
0 of 2 checks passed
@gdalle gdalle deleted the gd/batched branch June 20, 2024 08:58
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants