WIP: Add support for HyperHessian backend #940
Conversation
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #940 +/- ##
==========================================
- Coverage 98.23% 2.93% -95.30%
==========================================
Files 133 100 -33
Lines 7968 5582 -2386
==========================================
- Hits 7827 164 -7663
- Misses 141 5418 +5277
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
KristofferC
force-pushed
the
kc/hyperhessians
branch
from
9d3de40 to
6517285
Compare
|
Thanks for opening a PR! |
|
As far as knowing what to overload, maybe this page can help: https://juliadiff.org/DifferentiationInterface.jl/DifferentiationInterface/stable/dev/contributing/ |
|
Thanks for opening a PR! I did so, feel free to cancel the running tests on previous commit |
|
You can also comment out the whole |
|
At least the tests for the HyperHessian backend seems to pass. |
|
What's my next step? Do I need to add |
|
Maybe give me a few days to review this first? Also, it would be the first backend that implements only a few first-order operations, so I'm not sure how to document that |
Maybe the table of supported features is enough? I could add gradient, etc, but the hyperdual number would then degenerate to a normal dual number, and it would feel a bit silly to just re-implement things that are functionally identical to ForwardDiff. |
|
It's not a table of supported features, it's a table of which operators are implemented natively by the backend. For instance, Mooncake only implements a pullback, but thanks to DI's internal machinery you get a Jacobian operator for free. With this package, there is no pullback or pushforward so that machinery will fail |
|
It's not necessarily a bad thing, just a new thing |
gdalle
left a comment
gdalle
left a comment
There was a problem hiding this comment.
Overall this seems reasonable, the main decision to make now is whether you want to support Cache or not. I don't think it is crucial at first glance.
Could you maybe comment on #828 to see if I missed anything vis-a-vis creating a new backend? I should probably turn that issue into docs but I'd be curious about your (LLM-enhanced) experience
| be set to a positive `Int` to override HyperHessians' chunk heuristic; `nothing` | ||
| lets HyperHessians choose. | ||
| """ | ||
| struct AutoHyperHessians{CS} <: ADTypes.AbstractADType |
There was a problem hiding this comment.
Just adding a comment to remember we shouldn't merge this, since it belongs in ADTypes
| | `AutoFiniteDifferences` | ✅ | ✅ | | ||
| | `AutoForwardDiff` | ✅ | ✅ | | ||
| | `AutoGTPSA` | ✅ | ❌ | | ||
| | `AutoHyperHessians` | ✅ | ✅ | |
There was a problem hiding this comment.
I don't think you support caches, that would require dealing with hyperdual element types
| import DifferentiationInterface as DI | ||
| import .DI: AutoHyperHessians | ||
| using ADTypes: ForwardMode | ||
| using HyperHessians: |
There was a problem hiding this comment.
Are all of these public, as in, protected from API changes in minor or patch versions?
| ## Traits | ||
| DI.check_available(::DI.AutoHyperHessians) = true | ||
| DI.inplace_support(::DI.AutoHyperHessians) = DI.InPlaceSupported() | ||
| DI.hvp_mode(::DI.AutoHyperHessians) = DI.ForwardOverForward() |
There was a problem hiding this comment.
| DI.hvp_mode(::DI.AutoHyperHessians) = DI.ForwardOverForward() |
This is deduced automatically from the mode
| chunk_from_backend(backend::DI.AutoHyperHessians, N::Integer, ::Type{T}) where {T} = | ||
| isnothing(backend.chunksize) ? Chunk(pickchunksize(N, T), T) : Chunk{backend.chunksize}() | ||
|
|
||
| function DI.pick_batchsize(backend::DI.AutoHyperHessians, x::AbstractArray) |
There was a problem hiding this comment.
You're missing the same function for integer arguments instead of arrays
| contexts::Vararg{DI.Context, C}, | ||
| ) where {C} | ||
| DI.check_prep(f, prep, backend, x, contexts...) | ||
| fc = DI.fix_tail(f, map(DI.unwrap, contexts)...) |
There was a problem hiding this comment.
This will only work with Constant contexts (see the ForwardDiff extension for how I deal with Cache)
|
Also, the tests are failing |
2 participants
HyperHessian (https://github.com/KristofferC/HyperHessians.jl) is a forward mode AD package that specializes in second-order derivatives (Hessians). I initially made it like 4 years ago as an experiment to see what would happen if you made an AD package using hyper duals over nested duals for second order derivative but only recently freshened it up to be a proper package. The hyper dual number approach does have some extra flexibility over nested duals, which allows it to squeeze out some more performance over e.g. ForwardDiff:
Hessian:
ackleyackleyackleyackleyrosenbrock_1rosenbrock_1rosenbrock_1rosenbrock_1Hvp
rosenbrock_1rosenbrock_1rosenbrock_1rosenbrock_1ackleyackleyackleyackley(note that I am not claiming that it is impossible to improve ForwardDiff further w.r.t hessians, I am just saying that these are the numbers as of right now)
The package is not registered yet, but I thought I would open this for some feedback anyway. Note that some of the stuff in here was written with an LLM because figuring out exactly what to overload, etc, was tricky. Hopefully, it is a good starting point to iterate on. I also haven't checked that any overhead of DI vs raw HyperHessians is negligible yet.