Update gradient interface, support AbstractDifferentiation #90
Conversation
Codecov ReportAttention:
Additional details and impacted files@@ Coverage Diff @@
## master #90 +/- ##
==========================================
- Coverage 89.84% 89.03% -0.81%
==========================================
Files 21 20 -1
Lines 955 985 +30
==========================================
+ Hits 858 877 +19
- Misses 97 108 +11 ☔ View full report in Codecov by Sentry. |
docs/src/guide/getting_started.jl
Outdated
| @@ -70,7 +73,8 @@ solution, iterations = ffb(x0=ones(2), f=quadratic_cost, g=box_indicator) | |||
|
|
|||
| # We can verify the correctness of the solution by checking that the negative gradient is orthogonal to the constraints, pointing outwards: | |||
|
|
|||
| -ProximalAlgorithms.gradient(quadratic_cost, solution)[1] | |||
| v, pb = ProximalAlgorithms.value_and_pullback(quadratic_cost, solution) | |||
| -pb() | |||
There was a problem hiding this comment.
Maybe explain that the pullback outputs the gradient?
|
|
||
| function value_and_pullback(f::AutoDifferentiable, x) | ||
| fx, pb = AbstractDifferentiation.value_and_pullback_function(f.backend, f.f, x) | ||
| return fx, () -> pb(one(fx))[1] |
There was a problem hiding this comment.
Striclty speaking this is not a pullback if it takes no input. The point of a pullback is to take a cotangent and pull it back into the input space.
In this case, the cotangent is a scalar, and taking this scalar to be 1 returns the gradient, so it's okay, but the terminology might confuse autodiff people
There was a problem hiding this comment.
Yeah right, I was unsure here, thanks for pointing that out. I guess I could call it simply “closure”
There was a problem hiding this comment.
value_and_gradient_closure other something like that
src/ProximalAlgorithms.jl
Outdated
| """ | ||
| AutoDifferentiable(f, backend) | ||
|
|
||
| Wrap function `f` to be auto-differentiated using `backend`. |
There was a problem hiding this comment.
Specify that it's a callable struct?
|
|
||
| (f::ZygoteFunction)(x) = f.f(x) | ||
|
|
||
| function ProximalCore.gradient!(grad, f::ZygoteFunction, x) |
There was a problem hiding this comment.
The diff would be much smaller if you kept your functions gradient and gradient! but made them rely on value_and_pullback instead, right?
There was a problem hiding this comment.
I guess, but I’m thinking that these primitives don’t really fit the interface of ProximalCore maybe, see JuliaFirstOrder/ProximalCore.jl#3
Co-authored-by: Guillaume Dalle <[email protected]>
value_and_pullback(similar to AbstractDifferentiationvalue_and_pullback_function)Work in progress, opening for feedback and to take a look at CI workflows.
TODOs