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

Add support for Jacobians #22

Closed
wants to merge 14 commits into from
Closed

Conversation

adrhill
Copy link
Collaborator

@adrhill adrhill commented Mar 1, 2024

This PR adds fallback implementatios for value_and_jacobian!, value_and_jacobian, jacobian!, and jacobian.

The implementation dispatches on AbstractForwardBackend and AbstractReverseBackend:

  • AbstractForwardBackend: Makes use of value_and_pushforward!, which computes JVPs. The pushforward of a standard basis vector $e_i$ computes the $i$-th column of a Jacobian. The Jacobian is therefore assembled column-wise.
  • AbstractReverseBackend: Makes use of value_and_pullback!, which compute VJPs. The pullback of a standard basis vector $e_i$ computes the $i$-th row of a Jacobian. The Jacobian is therefore assembled row-wise.

This requires the implementation of a standard basis vector in src/unitvector.
This implementation can be replaced with more performant implementations in the future by specializing on backend types.

To write into pre-allocated Jacobians, type annotations of pullbacks and pushforwards had to be loosened.

Towards #6.

@codecov-commenter
Copy link

Codecov Report

Attention: Patch coverage is 85.07463% with 10 lines in your changes are missing coverage. Please review.

Project coverage is 88.82%. Comparing base (c65fb27) to head (538fb6b).

Files Patch % Lines
src/jacobian.jl 87.23% 6 Missing ⚠️
src/forward.jl 50.00% 2 Missing ⚠️
src/reverse.jl 50.00% 2 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main      #22      +/-   ##
==========================================
+ Coverage   85.98%   88.82%   +2.84%     
==========================================
  Files           9       11       +2     
  Lines         107      170      +63     
==========================================
+ Hits           92      151      +59     
- Misses         15       19       +4     

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

@adrhill
Copy link
Collaborator Author

adrhill commented Mar 2, 2024

Good news: only Enzyme required specialized unitvector and value_and_pullback! methods to support inplace mutations of Jacobians!

In both cases, this was due to Enzyme's Duplicated(x::T, dx::T) assuming identical types for inputs and outputs of the pullback.


One annoyance encountered in this PR is that value_and_pullback!(dx, ..., x::Real) assumes that dx is scalar and can't be mutated, even though it might in fact be a mutable entry in a pre-allocated Jacobian. A similar case exists for pushforwards on functions with scalar outputs y. This required an additional two methods for _value_and_jacobian!:
https://github.com/gdalle/DifferentiationInterface.jl/blob/1e40a0d01a76eb3c78141fffbad59bf9bf83add1/src/jacobian.jl#L41-L59

@adrhill adrhill marked this pull request as ready for review March 2, 2024 16:47
@gdalle gdalle mentioned this pull request Mar 6, 2024
@adrhill
Copy link
Collaborator Author

adrhill commented Mar 6, 2024

Superseded by #23.

@adrhill adrhill closed this Mar 6, 2024
@adrhill adrhill deleted the ah/jacobian branch March 6, 2024 18:51
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