More fixes

Author: devmotion

ext/AbstractDifferentiationChainRulesCoreExt.jl

@@ -1,13 +1,9 @@
 module AbstractDifferentiationChainRulesCoreExt
 
 import AbstractDifferentiation as AD
-if AD.EXTENSIONS_SUPPORTED
-    using ChainRulesCore: ChainRulesCore
-else
-    using .ChainRulesCore: ChainRulesCore
-end
+using ChainRulesCore: ChainRulesCore
 
-AD.@primitive function pullback_function(ab::AD.ReverseRuleConfigBackend, f, xs...)
+AD.@primitive function pullback_function(ba::AD.ReverseRuleConfigBackend, f, xs...)
     _, back = ChainRulesCore.rrule_via_ad(ab.ruleconfig, f, xs...)
     pullback(vs) = Base.tail(back(vs))
     pullback(vs::Tuple{Any}) = Base.tail(back(first(vs)))

ext/AbstractDifferentiationForwardDiffExt.jl

@@ -24,7 +24,7 @@ function AD.ForwardDiffBackend(; chunksize::Union{Val,Nothing}=nothing)
     return AD.ForwardDiffBackend{getchunksize(chunksize)}()
 end
 
-AD.@primitive function pushforward_function(::AD.ForwardDiffBackend, f, xs...)
+AD.@primitive function pushforward_function(ba::AD.ForwardDiffBackend, f, xs...)
     return function pushforward(vs)
         if length(xs) == 1
             v = vs isa Tuple ? only(vs) : vs

ext/AbstractDifferentiationReverseDiffExt.jl

@@ -13,7 +13,7 @@ AD.primal_value(x::ReverseDiff.TrackedReal) = ReverseDiff.value(x)
 AD.primal_value(x::AbstractArray{<:ReverseDiff.TrackedReal}) = ReverseDiff.value.(x)
 AD.primal_value(x::ReverseDiff.TrackedArray) = ReverseDiff.value(x)
 
-AD.@primitive function jacobian(::AD.ReverseDiffBackend, f, xs...)
+AD.@primitive function jacobian(ba::AD.ReverseDiffBackend, f, xs...)
     xs_arr = map(AD.asarray, xs)
     tape = ReverseDiff.JacobianTape(xs_arr) do (xs_arr...)
         xs_new = map(xs, xs_arr) do x, x_arr

ext/AbstractDifferentiationTrackerExt.jl

@@ -15,7 +15,7 @@ AD.primal_value(x::Tracker.TrackedReal) = Tracker.data(x)
 AD.primal_value(x::Tracker.TrackedArray) = Tracker.data(x)
 AD.primal_value(x::AbstractArray{<:Tracker.TrackedReal}) = Tracker.data.(x)
 
-AD.@primitive function pullback_function(::AD.TrackerBackend, f, xs...)
+AD.@primitive function pullback_function(ba::AD.TrackerBackend, f, xs...)
     value, back = Tracker.forward(f, xs...)
     function pullback(ws)
         if ws isa Tuple && !(value isa Tuple) 

src/AbstractDifferentiation.jl

@@ -506,16 +506,16 @@ macro primitive(expr)
 end
 
 function define_pushforward_function_and_friends(fdef)
-    fdef[:name] = :(AbstractDifferentiation.pushforward_function)
+    fdef[:name] = :($(AbstractDifferentiation).pushforward_function)
     args = fdef[:args]
     funcs = quote
         $(ExprTools.combinedef(fdef))
-        function AbstractDifferentiation.jacobian($(args...),)
-            identity_like = AbstractDifferentiation.identity_matrix_like($(args[3:end]...),)
-            pff = AbstractDifferentiation.pushforward_function($(args...),)
+        function $(AbstractDifferentiation).jacobian($(args...),)
+            identity_like = $(identity_matrix_like)($(args[3:end]...),)
+            pff = $(pushforward_function)($(args...),)
             if eltype(identity_like) <: Tuple{Vararg{Union{AbstractMatrix, Number}}}
                 return map(identity_like) do identity_like_i
-                    return mapreduce(hcat, AbstractDifferentiation._eachcol.(identity_like_i)...) do (cols...)
+                    return mapreduce(hcat, $(_eachcol).(identity_like_i)...) do (cols...)
                         pff(cols)
                     end
                 end
@@ -542,17 +542,17 @@ function define_pushforward_function_and_friends(fdef)
 end
 
 function define_pullback_function_and_friends(fdef)
-    fdef[:name] = :(AbstractDifferentiation.pullback_function)
+    fdef[:name] = :($(AbstractDifferentiation).pullback_function)
     args = fdef[:args]
     funcs = quote
         $(ExprTools.combinedef(fdef))
-        function AbstractDifferentiation.jacobian($(args...),)
-            value_and_pbf = AbstractDifferentiation.value_and_pullback_function($(args...),)
+        function $(AbstractDifferentiation).jacobian($(args...),)
+            value_and_pbf = $(value_and_pullback_function)($(args...),)
             value, _ = value_and_pbf(nothing)
-            identity_like = AbstractDifferentiation.identity_matrix_like(value)
+            identity_like = $(identity_matrix_like)(value)
             if eltype(identity_like) <: Tuple{Vararg{AbstractMatrix}}
                 return map(identity_like) do identity_like_i
-                    return mapreduce(vcat, AbstractDifferentiation._eachcol.(identity_like_i)...) do (cols...)
+                    return mapreduce(vcat, $(_eachcol).(identity_like_i)...) do (cols...)
                         value_and_pbf(cols)[2]'
                     end
                 end
@@ -575,12 +575,12 @@ _eachcol(a::Number) = (a,)
 _eachcol(a) = eachcol(a)
 
 function define_jacobian_and_friends(fdef)
-    fdef[:name] = :(AbstractDifferentiation.jacobian)
+    fdef[:name] = :($(AbstractDifferentiation).jacobian)
     return ExprTools.combinedef(fdef)
 end
 
 function define_primal_value(fdef)
-    fdef[:name] = :(AbstractDifferentiation.primal_value)
+    fdef[:name] = :($(AbstractDifferentiation).primal_value)
     return ExprTools.combinedef(fdef)
 end
 

src/backends.jl

@@ -54,7 +54,8 @@ struct TrackerBackend <: AbstractReverseMode end
 AD backend that uses reverse mode with any ChainRules-compatible reverse-mode AD package.
 
 !!! note
-    To be able to use this backend, you have to load ChainRulesCore.
+    On Julia >= 1.9, you have to load ChainRulesCore (possibly implicitly by loading
+    a ChainRules-compatible AD package) to be able to use this backend.
 """
 struct ReverseRuleConfigBackend{RC} <: AbstractReverseMode
     ruleconfig::RC

test/test_utils.jl

@@ -150,8 +150,8 @@ function test_hessians(backend; multiple_inputs=false, test_types=true)
     else
         # explicit test that AbstractDifferentiation throws an error
         # don't support tuple of Hessians
-        @test_throws AssertionError H1 = AD.hessian(backend, fgrad, (xvec, yvec))
-        @test_throws MethodError H1 = AD.hessian(backend, fgrad, xvec, yvec)
+        @test_throws ArgumentError AD.hessian(backend, fgrad, (xvec, yvec))
+        @test_throws MethodError AD.hessian(backend, fgrad, xvec, yvec)
     end
 
     # @test dfgraddxdx(xvec,yvec) ≈ H1[1] atol=1e-10