fixups and use ADTypes HyperHessians

Author: KristofferC

DifferentiationInterface/Project.toml

@@ -67,12 +67,12 @@ Diffractor = "=0.2.6"
 Enzyme = "0.13.39"
 EnzymeCore = "0.8.8"
 FastDifferentiation = "0.4.3"
-HyperHessians = "0.2"
 FiniteDiff = "2.27.0"
 FiniteDifferences = "0.12.31"
 ForwardDiff = "0.10.36,1"
 GPUArraysCore = "0.2"
 GTPSA = "1.4.0"
+HyperHessians = "0.2"
 LinearAlgebra = "1"
 Mooncake = "0.5.1 - 0.5.24"
 PolyesterForwardDiff = "0.1.2"

DifferentiationInterface/docs/src/explanation/backends.md

@@ -12,7 +12,7 @@ We support the following dense backend choices from [ADTypes.jl](https://github.
 - [`AutoFiniteDifferences`](@extref ADTypes.AutoFiniteDifferences)
 - [`AutoForwardDiff`](@extref ADTypes.AutoForwardDiff)
 - [`AutoGTPSA`](@extref ADTypes.AutoGTPSA)
-- [`AutoHyperHessians`](https://github.com/KristofferC/HyperHessians.jl)
+- [`AutoHyperHessians`](@extref ADTypes.AutoHyperHessians)
 - [`AutoMooncake`](@extref ADTypes.AutoMooncake) and [`AutoMooncakeForward`](@extref ADTypes.AutoMooncake) (the latter is experimental)
 - [`AutoPolyesterForwardDiff`](@extref ADTypes.AutoPolyesterForwardDiff)
 - [`AutoReverseDiff`](@extref ADTypes.AutoReverseDiff)
@@ -71,7 +71,7 @@ Moreover, each context type is supported by a specific subset of backends:
 | `AutoFiniteDifferences`    | ✅                  | ✅               |
 | `AutoForwardDiff`          | ✅                  | ✅               |
 | `AutoGTPSA`                | ✅                  | ❌               |
-| `AutoHyperHessians`        | ✅                  | ✅               |
+| `AutoHyperHessians`        | ✅                  | ❌               |
 | `AutoMooncake`             | ✅                  | ✅               |
 | `AutoMooncakeForward`      | ✅                  | ✅               |
 | `AutoPolyesterForwardDiff` | ✅                  | ✅               |

DifferentiationInterface/ext/DifferentiationInterfaceHyperHessiansExt/DifferentiationInterfaceHyperHessiansExt.jl

@@ -2,33 +2,29 @@ module DifferentiationInterfaceHyperHessiansExt
 
 import DifferentiationInterface as DI
 import .DI: AutoHyperHessians
-using ADTypes: ForwardMode
 using HyperHessians:
     HVPConfig,
     HessianConfig,
     Chunk,
     chunksize,
-    pickchunksize,
     HyperDual,
     hessian,
     hessian!,
     hessian_gradient_value,
     hessian_gradient_value!,
-    hessian,
     hvp,
     hvp!,
     hvp_gradient_value,
     hvp_gradient_value!
 
 ## Traits
 DI.check_available(::DI.AutoHyperHessians) = true
-DI.inplace_support(::DI.AutoHyperHessians) = DI.InPlaceSupported()
-DI.mode(::DI.AutoHyperHessians) = ForwardMode()
+DI.inplace_support(::DI.AutoHyperHessians) = DI.InPlaceNotSupported()
 
-chunk_from_backend(backend::DI.AutoHyperHessians, x) =
-    isnothing(backend.chunksize) ? Chunk(x) : Chunk{backend.chunksize}()
-chunk_from_backend(backend::DI.AutoHyperHessians, N::Integer, ::Type{T}) where {T} =
-    isnothing(backend.chunksize) ? Chunk(pickchunksize(N, T), T) : Chunk{backend.chunksize}()
+chunk_from_backend(::DI.AutoHyperHessians{nothing}, x) = Chunk(x)
+chunk_from_backend(::DI.AutoHyperHessians{CS}, x) where {CS} = Chunk{CS}()
+chunk_from_backend(::DI.AutoHyperHessians{nothing}, N::Integer, ::Type{T}) where {T} = Chunk(N, T)
+chunk_from_backend(::DI.AutoHyperHessians{CS}, ::Integer, ::Type) where {CS} = Chunk{CS}()
 
 function DI.pick_batchsize(backend::DI.AutoHyperHessians, x::AbstractArray)
     B = chunksize(chunk_from_backend(backend, x))
@@ -40,10 +36,11 @@ function DI.pick_batchsize(backend::DI.AutoHyperHessians, N::Integer)
     return DI.BatchSizeSettings{B}(N)
 end
 
-function DI.threshold_batchsize(backend::DI.AutoHyperHessians, chunksize2::Integer)
-    chunksize1 = backend.chunksize
-    chunksize = isnothing(chunksize1) ? nothing : min(chunksize1, chunksize2)
-    return DI.AutoHyperHessians(; chunksize)
+function DI.threshold_batchsize(::DI.AutoHyperHessians{nothing}, ::Integer)
+    return DI.AutoHyperHessians()
+end
+function DI.threshold_batchsize(::DI.AutoHyperHessians{CS}, chunksize2::Integer) where {CS}
+    return DI.AutoHyperHessians(; chunksize = min(CS, chunksize2))
 end
 
 function _translate_toprep(::Type{T}, c::Union{DI.GeneralizedConstant, DI.ConstantOrCache}) where {T}
@@ -111,8 +108,8 @@ function DI.second_derivative!(
         contexts::Vararg{DI.Context, C},
     ) where {C}
     DI.check_prep(f, prep, backend, x, contexts...)
-    copyto!(der2, DI.second_derivative(f, prep, backend, x, contexts...))
-    return der2
+    new_der2 = DI.second_derivative(f, prep, backend, x, contexts...)
+    return copyto!(der2, new_der2)
 end
 
 function DI.value_derivative_and_second_derivative(
@@ -253,7 +250,7 @@ end
 function DI.hvp(
         f,
         prep::HyperHessiansHVPPrep,
-        backend::AutoHyperHessians,
+        backend::DI.AutoHyperHessians,
         x,
         tx::NTuple,
         contexts::Vararg{DI.Context, C},

DifferentiationInterface/src/DifferentiationInterface.jl

@@ -33,7 +33,8 @@ using ADTypes:
     AutoReverseDiff,
     AutoSymbolics,
     AutoTracker,
-    AutoZygote
+    AutoZygote,
+    AutoHyperHessians
 using LinearAlgebra: dot
 
 include("compat.jl")
@@ -64,7 +65,6 @@ include("second_order/hessian.jl")
 
 include("misc/differentiate_with.jl")
 include("misc/from_primitive.jl")
-include("misc/autohyperhessians.jl")
 include("misc/sparsity_detector.jl")
 include("misc/simple_finite_diff.jl")
 include("misc/zero_backends.jl")

DifferentiationInterface/src/misc/autohyperhessians.jl

@@ -16,7 +16,7 @@ backends = [
 
 for backend in backends
     @test DI.check_available(backend)
-    @test DI.check_inplace(backend)
+    @test !DI.check_inplace(backend)
 end
 
 scenarios = default_scenarios(; include_constantified = true, include_cachified = true)