diff --git a/Project.toml b/Project.toml index ae0beacb8..39dfdfc80 100644 --- a/Project.toml +++ b/Project.toml @@ -3,11 +3,11 @@ uuid = "76274a88-744f-5084-9051-94815aaf08c4" version = "0.15.20" [deps] +ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197" ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" ChangesOfVariables = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0" -DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869" @@ -28,6 +28,7 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [weakdeps] ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2" DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c" +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" LazyArrays = "5078a376-72f3-5289-bfd5-ec5146d43c02" Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" @@ -35,6 +36,7 @@ ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" [extensions] BijectorsDistributionsADExt = "DistributionsAD" +BijectorsEnzymeExt = "Enzyme" BijectorsForwardDiffExt = "ForwardDiff" BijectorsLazyArraysExt = "LazyArrays" BijectorsMooncakeExt = "Mooncake" @@ -42,15 +44,16 @@ BijectorsReverseDiffChainRulesExt = ["ChainRules", "ReverseDiff"] BijectorsReverseDiffExt = "ReverseDiff" [compat] +ADTypes = "1" AbstractPPL = "0.14" ArgCheck = "1, 2" ChainRules = "1" ChainRulesCore = "0.10.11, 1" ChangesOfVariables = "0.1" -DifferentiationInterface = "0.7.14" Distributions = "0.25.33" DistributionsAD = "0.6" DocStringExtensions = "0.9" +Enzyme = "0.13" EnzymeCore = "0.8.15" FillArrays = "1" ForwardDiff = "0.10, 1.0.1" @@ -60,7 +63,7 @@ IrrationalConstants = "0.1, 0.2" LazyArrays = "2" LogExpFunctions = "0.3.3" MappedArrays = "0.2.2, 0.3, 0.4" -Mooncake = "0.4.95, 0.5" +Mooncake = "0.5.26" Reexport = "0.2, 1" ReverseDiff = "1" Roots = "1.3.15, 2" @@ -70,6 +73,7 @@ julia = "1.10.8" [extras] DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c" +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" LazyArrays = "5078a376-72f3-5289-bfd5-ec5146d43c02" diff --git a/docs/Project.toml b/docs/Project.toml index 9698442af..6e0d207a8 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -1,7 +1,6 @@ [deps] Bijectors = "76274a88-744f-5084-9051-94815aaf08c4" ChangesOfVariables = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0" -DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" diff --git a/docs/src/defining_examples.md b/docs/src/defining_examples.md index 95556b844..4c706853b 100644 --- a/docs/src/defining_examples.md +++ b/docs/src/defining_examples.md @@ -207,12 +207,10 @@ function full_transform(x12) return StereographicProj()(x123) end -import DifferentiationInterface as DI using FiniteDifferences, LinearAlgebra x = [0.3, 0.4, sgn * sqrt(1 - 0.3^2 - 0.4^2)] -adtype = DI.AutoFiniteDifferences(; fdm=central_fdm(5, 1)) -jac = DI.jacobian(full_transform, adtype, x[1:2]) +jac = only(FiniteDifferences.jacobian(central_fdm(5, 1), full_transform, x[1:2])) logjac = logabsdet(jac)[1] ``` diff --git a/ext/BijectorsEnzymeExt.jl b/ext/BijectorsEnzymeExt.jl new file mode 100644 index 000000000..af1ce5638 --- /dev/null +++ b/ext/BijectorsEnzymeExt.jl @@ -0,0 +1,105 @@ +module BijectorsEnzymeExt + +import Bijectors: _value_and_gradient, _value_and_jacobian +import ADTypes: AutoEnzyme +using Enzyme: Enzyme +using EnzymeCore: EnzymeCore + +const DuplicatedFunctionAnnotations = Union{ + EnzymeCore.Duplicated,EnzymeCore.DuplicatedNoNeed,EnzymeCore.MixedDuplicated +} + +function _annotate_function(f, backend::AutoEnzyme, mode) + annotation = typeof(backend).parameters[2] + if annotation === Nothing + return f + elseif annotation <: EnzymeCore.Const + return Enzyme.Const(f) + elseif annotation <: DuplicatedFunctionAnnotations + if Enzyme.guess_activity(typeof(f), mode) <: EnzymeCore.Const + return Enzyme.Const(f) + else + # Enzyme's sugar APIs only preserve function shadows for `Duplicated`, + # so normalize the duplicated-like annotations here. + return Enzyme.Duplicated(f, Enzyme.make_zero(f)) + end + else + throw(ArgumentError("unsupported Enzyme function annotation $annotation")) + end +end + +function _value_and_gradient( + f, + backend::Union{AutoEnzyme{Nothing},AutoEnzyme{<:EnzymeCore.ReverseMode}}, + x::AbstractVector, +) + mode = if backend isa AutoEnzyme{Nothing} + Enzyme.ReverseWithPrimal + else + Enzyme.WithPrimal(backend.mode) + end + annotated_f = _annotate_function(f, backend, mode) + dx = zero(x) + _, val = Enzyme.autodiff(mode, annotated_f, Enzyme.Active, Enzyme.Duplicated(x, dx)) + return val, dx +end + +function _value_and_gradient( + f, backend::AutoEnzyme{<:EnzymeCore.ForwardMode}, x::AbstractVector +) + mode = Enzyme.WithPrimal(backend.mode) + annotated_f = _annotate_function(f, backend, mode) + grad = zero(x) + value = f(x) + for i in eachindex(x) + dx = zero(x) + dx[i] = one(eltype(x)) + directional, primal = Enzyme.autodiff(mode, annotated_f, Enzyme.Duplicated(x, dx)) + grad[i] = directional + if i == firstindex(x) + value = primal + end + end + return value, grad +end + +function _value_and_jacobian( + f, backend::AutoEnzyme{<:EnzymeCore.ReverseMode}, x::AbstractVector +) + value = f(x) + if isempty(x) || isempty(value) + return value, Matrix{eltype(value)}(undef, length(value), length(x)) + end + annotated_f = _annotate_function(f, backend, backend.mode) + jacobian = only(Enzyme.jacobian(backend.mode, annotated_f, x)) + return value, reshape(jacobian, length(value), length(x)) +end + +function _value_and_jacobian(f, ::AutoEnzyme{Nothing}, x::AbstractVector) + return _value_and_jacobian(f, AutoEnzyme(; mode=Enzyme.Forward), x) +end + +function _value_and_jacobian( + f, backend::AutoEnzyme{<:EnzymeCore.ForwardMode}, x::AbstractVector +) + mode = Enzyme.WithPrimal(backend.mode) + annotated_f = _annotate_function(f, backend, mode) + value = f(x) + J = nothing + for i in eachindex(x) + dx = zero(x) + dx[i] = one(eltype(x)) + directional, primal = Enzyme.autodiff(mode, annotated_f, Enzyme.Duplicated(x, dx)) + if i == firstindex(x) + value = primal isa AbstractArray ? copy(primal) : primal + J = Matrix{eltype(directional)}(undef, length(directional), length(x)) + end + J[:, i] .= directional + end + if isnothing(J) + J = Matrix{eltype(value)}(undef, length(value), 0) + end + return value, J +end + +end diff --git a/ext/BijectorsForwardDiffExt.jl b/ext/BijectorsForwardDiffExt.jl index 76ae0dd2f..4a4f9f27c 100644 --- a/ext/BijectorsForwardDiffExt.jl +++ b/ext/BijectorsForwardDiffExt.jl @@ -1,8 +1,44 @@ module BijectorsForwardDiffExt -using Bijectors: Bijectors, find_alpha +import Bijectors: Bijectors, find_alpha, _value_and_gradient, _value_and_jacobian +import ADTypes: AutoForwardDiff using ForwardDiff: ForwardDiff +function _value_and_gradient( + f, backend::AutoForwardDiff{chunksize,T}, x::AbstractVector +) where {chunksize,T} + if isempty(x) + return f(x), similar(x, 0) + end + result = ForwardDiff.DiffResults.GradientResult(x) + chunk = isnothing(chunksize) ? ForwardDiff.Chunk(x) : ForwardDiff.Chunk{chunksize}() + tag = T === Nothing ? ForwardDiff.Tag(f, eltype(x)) : backend.tag + config = ForwardDiff.GradientConfig(nothing, x, chunk, tag) + if T === Nothing + ForwardDiff.checktag(config, f, x) + end + ForwardDiff.gradient!(result, f, x, config, Val(false)) + return ForwardDiff.DiffResults.value(result), ForwardDiff.DiffResults.gradient(result) +end + +function _value_and_jacobian( + f, backend::AutoForwardDiff{chunksize,T}, x::AbstractVector +) where {chunksize,T} + y = f(x) + if isempty(x) + return y, Matrix{eltype(y)}(undef, length(y), 0) + end + result = ForwardDiff.DiffResults.JacobianResult(y, x) + chunk = isnothing(chunksize) ? ForwardDiff.Chunk(x) : ForwardDiff.Chunk{chunksize}() + tag = T === Nothing ? ForwardDiff.Tag(f, eltype(x)) : backend.tag + config = ForwardDiff.JacobianConfig(nothing, x, chunk, tag) + if T === Nothing + ForwardDiff.checktag(config, f, x) + end + ForwardDiff.jacobian!(result, f, x, config, Val(false)) + return ForwardDiff.DiffResults.value(result), ForwardDiff.DiffResults.jacobian(result) +end + Bijectors._eps(::Type{<:ForwardDiff.Dual{<:Any,Real}}) = Bijectors._eps(Real) Bijectors._eps(::Type{<:ForwardDiff.Dual{<:Any,<:Integer}}) = Bijectors._eps(Real) diff --git a/ext/BijectorsLazyArraysExt.jl b/ext/BijectorsLazyArraysExt.jl index 03943fdc7..6be9502a5 100644 --- a/ext/BijectorsLazyArraysExt.jl +++ b/ext/BijectorsLazyArraysExt.jl @@ -1,15 +1,15 @@ module BijectorsLazyArraysExt import Bijectors: maporbroadcast -using LazyArrays: LazyArrays +using LazyArrays: BroadcastArray -function maporbroadcast(f, x1::LazyArrays.BroadcastArray, x...) +function maporbroadcast(f, x1::BroadcastArray, x...) return copy(f.(x1, x...)) end -function maporbroadcast(f, x1, x2::LazyArrays.BroadcastArray, x...) +function maporbroadcast(f, x1, x2::BroadcastArray, x...) return copy(f.(x1, x2, x...)) end -function maporbroadcast(f, x1, x2, x3::LazyArrays.BroadcastArray, x...) +function maporbroadcast(f, x1, x2, x3::BroadcastArray, x...) return copy(f.(x1, x2, x3, x...)) end diff --git a/ext/BijectorsMooncakeExt.jl b/ext/BijectorsMooncakeExt.jl index bea5c2433..6f2af0216 100644 --- a/ext/BijectorsMooncakeExt.jl +++ b/ext/BijectorsMooncakeExt.jl @@ -1,8 +1,101 @@ module BijectorsMooncakeExt +using Mooncake: Mooncake using Mooncake: - @is_primitive, MinimalCtx, Mooncake, CoDual, primal, tangent_type, @from_chainrules + @is_primitive, + MinimalCtx, + CoDual, + tangent_type, + @from_chainrules, + prepare_pullback_cache, + prepare_gradient_cache, + prepare_derivative_cache, + value_and_pullback!!, + value_and_gradient!!, + value_and_derivative!!, + zero_tangent, + Config, + _copy_output, + tangent_to_primal!! using Bijectors: find_alpha, ChainRulesCore +import Bijectors: _value_and_gradient, _value_and_jacobian +import ADTypes: AutoMooncake, AutoMooncakeForward + +_mooncake_config(::Union{AutoMooncake{Nothing},AutoMooncakeForward{Nothing}}) = Config() +_mooncake_config(backend::Union{AutoMooncake,AutoMooncakeForward}) = backend.config + +function _mooncake_zero_tangent_or_primal( + x, backend::Union{AutoMooncake,AutoMooncakeForward} +) + if _mooncake_config(backend).friendly_tangents + return tangent_to_primal!!(_copy_output(x), zero_tangent(x)) + else + return zero_tangent(x) + end +end + +## Reverse-mode implementations + +function _value_and_gradient(f, backend::AutoMooncake, x::AbstractVector) + cache = prepare_gradient_cache(f, x; config=_mooncake_config(backend)) + val, (_, x_grad) = value_and_gradient!!(cache, f, x) + return val, x_grad +end + +function _value_and_jacobian(f, backend::AutoMooncake, x::AbstractVector) + cache = prepare_pullback_cache(f, x; config=_mooncake_config(backend)) + n_out, n_in = length(cache.y_cache), length(x) + dy = zeros(eltype(cache.y_cache), n_out) + if n_out > 0 + dy[1] = one(eltype(cache.y_cache)) + end + val, (_, first_row) = value_and_pullback!!(cache, dy, f, x) + if n_out == 0 + return _copy_output(val), Matrix{eltype(x)}(undef, 0, n_in) + end + y = _copy_output(val) + J = Matrix{eltype(first_row)}(undef, n_out, n_in) + J[1, :] .= first_row + for i in 2:n_out + fill!(dy, zero(eltype(cache.y_cache))) + dy[i] = one(eltype(cache.y_cache)) + _, (_, row) = value_and_pullback!!(cache, dy, f, x) + J[i, :] .= row + end + return y, J +end + +## Forward-mode implementations (column-by-column JVPs) + +function _value_and_gradient(f, backend::AutoMooncakeForward, x::AbstractVector) + cache = prepare_gradient_cache(f, x; config=_mooncake_config(backend)) + val, (_, x_grad) = value_and_gradient!!(cache, f, x) + return val, x_grad +end + +function _value_and_jacobian( + f, backend::AutoMooncakeForward, x::AbstractVector{T} +) where {T} + y = f(x) + n_out, n_in = length(y), length(x) + cache = prepare_derivative_cache(f, x; config=_mooncake_config(backend)) + df = _mooncake_zero_tangent_or_primal(f, backend) + if n_in == 0 + return y, Matrix{eltype(y)}(undef, n_out, 0) + end + dx = zeros(T, n_in) + dx[1] = one(T) + _, first_jvp = value_and_derivative!!(cache, (f, df), (x, dx)) + J = Matrix{eltype(first_jvp)}(undef, n_out, n_in) + J[:, 1] .= first_jvp + for j in 2:n_in + fill!(dx, zero(T)) + dx[j] = one(T) + _, jvp = value_and_derivative!!(cache, (f, df), (x, dx)) + J[:, j] .= jvp + end + return y, J +end @from_chainrules(MinimalCtx, Tuple{typeof(find_alpha),Float16,Float16,Float16}) @from_chainrules(MinimalCtx, Tuple{typeof(find_alpha),Float32,Float32,Float32}) @@ -50,7 +143,9 @@ function Mooncake.rrule!!( msg = "Integer argument has tangent type $(tangent_type(I)), should be NoTangent." throw(ArgumentError(msg)) end - out, pb = ChainRulesCore.rrule(find_alpha, primal(x), primal(y), primal(z)) + out, pb = ChainRulesCore.rrule( + find_alpha, Mooncake.primal(x), Mooncake.primal(y), Mooncake.primal(z) + ) function find_alpha_pb(dout::P) _, dx, dy, _ = pb(dout) return Mooncake.NoRData(), P(dx), P(dy), Mooncake.NoRData() diff --git a/ext/BijectorsReverseDiffExt.jl b/ext/BijectorsReverseDiffExt.jl index ea5b959e2..1a54d2c96 100644 --- a/ext/BijectorsReverseDiffExt.jl +++ b/ext/BijectorsReverseDiffExt.jl @@ -10,15 +10,52 @@ using ReverseDiff: TrackedMatrix, @grad_from_chainrules +import Bijectors: _value_and_gradient, _value_and_jacobian +import ADTypes: AutoReverseDiff + +function _value_and_gradient(f, ::AutoReverseDiff{false}, x::AbstractVector) + if isempty(x) + return f(x), similar(x, 0) + end + result = ReverseDiff.DiffResults.GradientResult(x) + ReverseDiff.gradient!(result, f, x) + return ReverseDiff.DiffResults.value(result), ReverseDiff.DiffResults.gradient(result) +end + +function _value_and_jacobian(f, ::AutoReverseDiff{false}, x::AbstractVector) + y = f(x) + result = ReverseDiff.DiffResults.JacobianResult(y, x) + ReverseDiff.jacobian!(result, f, x) + return ReverseDiff.DiffResults.value(result), ReverseDiff.DiffResults.jacobian(result) +end + +function _value_and_gradient(f, ::AutoReverseDiff{true}, x::AbstractVector) + if isempty(x) + return f(x), similar(x, 0) + end + tape = ReverseDiff.GradientTape(f, x) + compiled = ReverseDiff.compile(tape) + result = ReverseDiff.DiffResults.GradientResult(x) + ReverseDiff.gradient!(result, compiled, x) + return ReverseDiff.DiffResults.value(result), ReverseDiff.DiffResults.gradient(result) +end + +function _value_and_jacobian(f, ::AutoReverseDiff{true}, x::AbstractVector) + tape = ReverseDiff.JacobianTape(f, x) + compiled = ReverseDiff.compile(tape) + y = f(x) + result = ReverseDiff.DiffResults.JacobianResult(y, x) + ReverseDiff.jacobian!(result, compiled, x) + return ReverseDiff.DiffResults.value(result), ReverseDiff.DiffResults.jacobian(result) +end + using Bijectors: - ChainRulesCore, Elementwise, SimplexBijector, maphcat, simplex_link_jacobian, simplex_invlink_jacobian, - simplex_logabsdetjac_gradient, - Inverse + simplex_logabsdetjac_gradient import Bijectors: Bijectors, _eps, @@ -27,8 +64,6 @@ import Bijectors: _simplex_bijector, _simplex_inv_bijector, replace_diag, - jacobian, - _inv_link_chol_lkj, _link_chol_lkj, _transform_ordered, _transform_inverse_ordered, diff --git a/src/Bijectors.jl b/src/Bijectors.jl index 3fd0ddba5..5c50692db 100644 --- a/src/Bijectors.jl +++ b/src/Bijectors.jl @@ -40,6 +40,7 @@ using InverseFunctions: InverseFunctions import ChangesOfVariables: ChangesOfVariables, with_logabsdet_jacobian import InverseFunctions: inverse +using ADTypes: ADTypes using ChainRulesCore: ChainRulesCore using Functors: Functors using IrrationalConstants: IrrationalConstants @@ -359,6 +360,7 @@ end include("utils.jl") include("interface.jl") include("chainrules.jl") +include("ad_utils.jl") include("vector/VectorBijectors.jl") # Broadcasting here breaks some AD backends for certain array types diff --git a/src/ad_utils.jl b/src/ad_utils.jl new file mode 100644 index 000000000..a4982005c --- /dev/null +++ b/src/ad_utils.jl @@ -0,0 +1,30 @@ +""" + _value_and_gradient(f, backend::ADTypes.AbstractADType, x::AbstractVector) + +Compute the value and gradient of scalar-valued `f` at `x`. +Returns `(f(x), gradient)` where `gradient` is a vector of the same size as `x`. + +Implementations are provided by package extensions for each AD backend. +""" +function _value_and_gradient(f, backend::ADTypes.AbstractADType, x::AbstractVector) + return error( + "_value_and_gradient is not implemented for backend $(typeof(backend)). " * + "Load the corresponding AD package to enable support.", + ) +end + +""" + _value_and_jacobian(f, backend::ADTypes.AbstractADType, x::AbstractVector) + +Compute the value and Jacobian of `f` at `x`. +Returns `(f(x), jacobian)` where `jacobian` is a matrix of size +`(length(f(x)), length(x))`. + +Implementations are provided by package extensions for each AD backend. +""" +function _value_and_jacobian(f, backend::ADTypes.AbstractADType, x::AbstractVector) + return error( + "_value_and_jacobian is not implemented for backend $(typeof(backend)). " * + "Load the corresponding AD package to enable support.", + ) +end diff --git a/src/bijectors/named_bijector.jl b/src/bijectors/named_bijector.jl index 85e53e9bc..4838a177d 100644 --- a/src/bijectors/named_bijector.jl +++ b/src/bijectors/named_bijector.jl @@ -139,7 +139,7 @@ deps(b::NamedCoupling{<:Any,Deps}) where {Deps} = Deps return quote b = nc.f($([:(x.$d) for d in deps]...)) x_target, logjac = with_logabsdet_jacobian(b, x.$target) - return merge(x, ($target=x_target,)), logjac + return merge(x, (($target)=x_target,)), logjac end end @@ -149,6 +149,6 @@ end return quote ib = inverse(ni.orig.f($([:(x.$d) for d in deps]...))) x_target, logjac = with_logabsdet_jacobian(ib, x.$target) - return merge(x, ($target=x_target,)), logjac + return merge(x, (($target)=x_target,)), logjac end end diff --git a/src/bijectors/named_stacked.jl b/src/bijectors/named_stacked.jl index 59076ac88..67d790b77 100644 --- a/src/bijectors/named_stacked.jl +++ b/src/bijectors/named_stacked.jl @@ -87,8 +87,8 @@ end end ), ) - push!(exprs, :(transforms = merge(transforms, ($n=trf,)))) - push!(exprs, :(ranges = merge(ranges, ($n=output_range,)))) + push!(exprs, :(transforms = merge(transforms, (($n)=trf,)))) + push!(exprs, :(ranges = merge(ranges, (($n)=output_range,)))) end push!(exprs, :(return NamedStacked{names}(transforms, ranges))) return Expr(:block, exprs...) @@ -153,14 +153,15 @@ end if i == 1 push!( exprs, - :(output = ($n=inverse(nsi.orig.transforms.$n)(y[nsi.orig.ranges.$n]),)), + :(output = (($n)=inverse(nsi.orig.transforms.$n)(y[nsi.orig.ranges.$n]),)), ) else push!( exprs, :( output = merge( - output, ($n=inverse(nsi.orig.transforms.$n)(y[nsi.orig.ranges.$n]),) + output, + (($n)=inverse(nsi.orig.transforms.$n)(y[nsi.orig.ranges.$n]),), ) ), ) @@ -182,7 +183,7 @@ end first_out, first_logjac = with_logabsdet_jacobian( inverse(nsi.orig.transforms.$n), y[nsi.orig.ranges.$n] ) - output = ($n=first_out,) + output = (($n)=first_out,) logjac = first_logjac end, ) @@ -193,7 +194,7 @@ end next_out, next_logjac = with_logabsdet_jacobian( inverse(nsi.orig.transforms.$n), y[nsi.orig.ranges.$n] ) - output = merge(output, ($n=next_out,)) + output = merge(output, (($n)=next_out,)) logjac += next_logjac end, ) diff --git a/src/vector/product/product.jl b/src/vector/product/product.jl index 33c784a3b..dc35b39e9 100644 --- a/src/vector/product/product.jl +++ b/src/vector/product/product.jl @@ -303,7 +303,7 @@ end expr.args, :((xr, lj) = with_logabsdet_jacobian(t.transforms.$nm, view(y, t.ranges.$nm))), ) - push!(expr.args, :(x = (x..., $nm=xr))) + push!(expr.args, :(x = (x..., ($nm)=xr))) push!(expr.args, :(logjac += lj)) end push!(expr.args, :(return x, logjac)) @@ -332,7 +332,7 @@ end push!(exprs, :(offset = 1)) for nm in names push!(exprs, :(this_length = length_fn(dists.$nm))) - push!(exprs, :(ranges = (ranges..., $nm=offset:(offset + this_length - 1)))) + push!(exprs, :(ranges = (ranges..., ($nm)=offset:(offset + this_length - 1)))) push!(exprs, :(offset += this_length)) end push!(exprs, :(return struct_type(trfms, ranges, size(dists[1])))) diff --git a/src/vector/test_utils.jl b/src/vector/test_utils.jl index bb5003f49..273f9d976 100644 --- a/src/vector/test_utils.jl +++ b/src/vector/test_utils.jl @@ -1,19 +1,25 @@ using Test using LinearAlgebra: logabsdet, Cholesky, UpperTriangular, LowerTriangular -import DifferentiationInterface as DI +import ADTypes: + AutoForwardDiff, + AutoReverseDiff, + AutoMooncake, + AutoMooncakeForward, + AutoEnzyme, + AbstractADType import EnzymeCore as EC # Would like to use FiniteDifferences, but very easy to run into issues with # https://juliadiff.org/FiniteDifferences.jl/latest/#Dealing-with-Singularities -const ref_adtype = DI.AutoForwardDiff() - -const default_adtypes = [ - DI.AutoReverseDiff(), - DI.AutoReverseDiff(; compile=true), - DI.AutoMooncake(), - DI.AutoMooncakeForward(), - DI.AutoEnzyme(; mode=EC.Forward, function_annotation=EC.Const), - DI.AutoEnzyme(; mode=EC.Reverse, function_annotation=EC.Const), +const ref_adtype = AutoForwardDiff() + +const default_adtypes = AbstractADType[ + AutoReverseDiff(), + AutoReverseDiff(; compile=true), + AutoMooncake(; config=nothing), + AutoMooncakeForward(; config=nothing), + AutoEnzyme(; mode=EC.Forward, function_annotation=EC.Const), + AutoEnzyme(; mode=EC.Reverse, function_annotation=EC.Const), ] _get_value_support(::D.Distribution{<:Any,VS}) where {VS<:D.ValueSupport} = VS @@ -422,7 +428,7 @@ function test_optics(d::D.Distribution) x = rand(d) xvec = to_vec(d)(x) yvec = to_linked_vec(d)(x) - J = DI.jacobian(to_linked_vec(d) ∘ from_vec(d), ref_adtype, xvec) + _, J = B._value_and_jacobian(to_linked_vec(d) ∘ from_vec(d), ref_adtype, xvec) o = optic_vec(d) lo = linked_optic_vec(d) for i in 1:length(yvec) @@ -566,7 +572,9 @@ function test_logjac(d::D.Distribution, atol, rtol) # to make sure that the Jacobian is square. ad_xvec = to_vec_for_logjac_test(d)(x) ad_ffwd = to_linked_vec(d) ∘ from_vec_for_logjac_test(d) - ad_logjac = first(logabsdet(DI.jacobian(ad_ffwd, ref_adtype, ad_xvec))) + ad_logjac = first( + logabsdet(last(B._value_and_jacobian(ad_ffwd, ref_adtype, ad_xvec))) + ) @test vbt_logjac ≈ ad_logjac atol = atol rtol = rtol end @@ -579,7 +587,9 @@ function test_logjac(d::D.Distribution, atol, rtol) # For the AD calculation we need to use to/from_vec_for_logjac_test instead, # to make sure that the Jacobian is square. ad_frvs = to_vec_for_logjac_test(d) ∘ from_linked_vec(d) - ad_logjac = first(logabsdet(DI.jacobian(ad_frvs, ref_adtype, yvec))) + ad_logjac = first( + logabsdet(last(B._value_and_jacobian(ad_frvs, ref_adtype, yvec))) + ) @test vbt_logjac ≈ ad_logjac atol = atol rtol = rtol end end @@ -590,7 +600,7 @@ end Test that various AD backends can differentiate the conversions to and from vector and linked vector forms for the given distribution `d`. """ -function test_ad(d::D.Distribution, adtypes::Vector{<:DI.AbstractADType}, atol, rtol) +function test_ad(d::D.Distribution, adtypes::Vector{<:AbstractADType}, atol, rtol) # If `d` is a discrete distribution, Mooncake refuses to differentiate through the # transforms (which are just identity transforms). Likewise, Enzyme will throw an # error saying that the output is Const but was not marked as such. @@ -600,9 +610,9 @@ function test_ad(d::D.Distribution, adtypes::Vector{<:DI.AbstractADType}, atol, adtypes = if d isa D.Distribution{<:Any,D.Discrete} filter(adtypes) do adtype !( - adtype isa DI.AutoMooncake || - adtype isa DI.AutoMooncakeForward || - adtype isa DI.AutoEnzyme + adtype isa AutoMooncake || + adtype isa AutoMooncakeForward || + adtype isa AutoEnzyme ) end else @@ -613,18 +623,18 @@ function test_ad(d::D.Distribution, adtypes::Vector{<:DI.AbstractADType}, atol, x = _rand_safe_ad(d) xvec = to_vec(d)(x) ffwd = to_linked_vec(d) ∘ from_vec(d) - ref_jac = DI.jacobian(ffwd, ref_adtype, xvec) + _, ref_jac = B._value_and_jacobian(ffwd, ref_adtype, xvec) ladj(xvec) = last(with_logabsdet_jacobian(ffwd, xvec)) - ref_grad_ladj = DI.gradient(ladj, ref_adtype, xvec) + _, ref_grad_ladj = B._value_and_gradient(ladj, ref_adtype, xvec) for adtype in adtypes @testset let x = x, adtype = adtype, d = d - ad_jac = DI.jacobian(ffwd, adtype, xvec) + _, ad_jac = B._value_and_jacobian(ffwd, adtype, xvec) @test ref_jac ≈ ad_jac atol = atol rtol = rtol end @testset let x = x, adtype = adtype, d = d - ad_grad_ladj = DI.gradient(ladj, adtype, xvec) + _, ad_grad_ladj = B._value_and_gradient(ladj, adtype, xvec) @test ref_grad_ladj ≈ ad_grad_ladj atol = atol rtol = rtol end end @@ -634,19 +644,19 @@ function test_ad(d::D.Distribution, adtypes::Vector{<:DI.AbstractADType}, atol, x = _rand_safe_ad(d) yvec = to_linked_vec(d)(x) frvs = to_vec(d) ∘ from_linked_vec(d) - ref_jac = DI.jacobian(frvs, ref_adtype, yvec) + _, ref_jac = B._value_and_jacobian(frvs, ref_adtype, yvec) ladj(yvec) = last(with_logabsdet_jacobian(frvs, yvec)) - ref_grad_ladj = DI.gradient(ladj, ref_adtype, yvec) + _, ref_grad_ladj = B._value_and_gradient(ladj, ref_adtype, yvec) for adtype in adtypes @testset let x = x, adtype = adtype, d = d - ad_jac = DI.jacobian(frvs, adtype, yvec) + _, ad_jac = B._value_and_jacobian(frvs, adtype, yvec) @test ref_jac ≈ ad_jac atol = atol rtol = rtol end @testset let x = x, adtype = adtype, d = d - ad_grad_ladj = DI.gradient(ladj, adtype, yvec) + _, ad_grad_ladj = B._value_and_gradient(ladj, adtype, yvec) @test ref_grad_ladj ≈ ad_grad_ladj atol = atol rtol = rtol end end diff --git a/test/Project.toml b/test/Project.toml index 8f6b34696..5c7a0a872 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,12 +1,13 @@ [deps] +ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001" AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" AdvancedHMC = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d" +Bijectors = "76274a88-744f-5084-9051-94815aaf08c4" ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2" ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a" ChangesOfVariables = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0" Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa" -DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" @@ -31,6 +32,7 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" [compat] +ADTypes = "1" AbstractMCMC = "5" AbstractPPL = "0.14" AdvancedHMC = "0.6, 0.7, 0.8" @@ -38,7 +40,6 @@ ChainRules = "1" ChainRulesTestUtils = "0.7, 1" ChangesOfVariables = "0.1" Combinatorics = "1.0.2" -DifferentiationInterface = "0.7.7" Distributions = "0.25" DistributionsAD = "0.6.3" Documenter = "1" @@ -52,7 +53,7 @@ LazyArrays = "1, 2" LogDensityProblems = "2" LogExpFunctions = "0.3.1" MCMCDiagnosticTools = "0.3" -Mooncake = "0.4, 0.5" +Mooncake = "0.5.26" PDMats = "0.11" ReverseDiff = "1.4.2" StableRNGs = "1" diff --git a/test/bijectors/product_bijector.jl b/test/bijectors/product_bijector.jl index 818f89c09..73dbba086 100644 --- a/test/bijectors/product_bijector.jl +++ b/test/bijectors/product_bijector.jl @@ -39,7 +39,7 @@ has_square_jacobian(b, x) = Bijectors.output_size(b, x) == size(x) y, logjac, changes_of_variables_test=has_square_jacobian(b, xs[1]), - test_not_identity=!isidentity, + test_not_identity=(!isidentity), ) end @@ -63,7 +63,7 @@ has_square_jacobian(b, x) = Bijectors.output_size(b, x) == size(x) y, logjac, changes_of_variables_test=has_square_jacobian(b, xs[1]), - test_not_identity=!isidentity, + test_not_identity=(!isidentity), ) end end diff --git a/test/bijectors/rational_quadratic_spline.jl b/test/bijectors/rational_quadratic_spline.jl index a936c01c8..cfb315d81 100644 --- a/test/bijectors/rational_quadratic_spline.jl +++ b/test/bijectors/rational_quadratic_spline.jl @@ -46,7 +46,7 @@ using LogExpFunctions # Outside of domain x = 5.0 - test_bijector(b, -x; y=-x, logjac=0) + test_bijector(b, -x; y=(-x), logjac=0) test_bijector(b, x; y=x, logjac=0) # multivariate diff --git a/test/integration/enzyme/Project.toml b/test/integration/enzyme/Project.toml index 3ce74b7ab..fe9714f77 100644 --- a/test/integration/enzyme/Project.toml +++ b/test/integration/enzyme/Project.toml @@ -1,7 +1,6 @@ [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" Bijectors = "76274a88-744f-5084-9051-94815aaf08c4" -DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" diff --git a/test/integration/enzyme/main.jl b/test/integration/enzyme/main.jl index 45cb5c5be..45a414a1a 100644 --- a/test/integration/enzyme/main.jl +++ b/test/integration/enzyme/main.jl @@ -1,17 +1,16 @@ using ADTypes using Bijectors -using DifferentiationInterface using Enzyme using FiniteDifferences using LinearAlgebra using Test -const REF_BACKEND = AutoFiniteDifferences(; fdm=central_fdm(5, 1)) +const REF_FDM = central_fdm(5, 1) function test_ad(f, backend, x; rtol=1e-6, atol=1e-6) @info "testing AD for function $f with $backend" - ref_gradient = DifferentiationInterface.gradient(f, REF_BACKEND, x) - gradient = DifferentiationInterface.gradient(f, backend, x) + ref_gradient = only(FiniteDifferences.grad(REF_FDM, f, x)) + _, gradient = Bijectors._value_and_gradient(f, backend, x) @test isapprox(gradient, ref_gradient; rtol=rtol, atol=atol) end diff --git a/test/runtests.jl b/test/runtests.jl index 97074c797..06fa45dcc 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,8 +1,8 @@ using Bijectors +using ADTypes using ChainRulesTestUtils using Combinatorics -using DifferentiationInterface using DistributionsAD using Documenter: Documenter using FiniteDifferences @@ -57,12 +57,12 @@ if TEST_ENZYME ), ] end -const REF_BACKEND = AutoFiniteDifferences(; fdm=central_fdm(5, 1)) +const REF_FDM = central_fdm(5, 1) function test_ad(f, backend, x; rtol=1e-6, atol=1e-6) @info "testing AD for function $f with $backend" - ref_gradient = DifferentiationInterface.gradient(f, REF_BACKEND, x) - gradient = DifferentiationInterface.gradient(f, backend, x) + ref_gradient = only(FiniteDifferences.grad(REF_FDM, f, x)) + _, gradient = Bijectors._value_and_gradient(f, backend, x) @test isapprox(gradient, ref_gradient; rtol=rtol, atol=atol) end diff --git a/test/vector/cholesky.jl b/test/vector/cholesky.jl index 74ea1fd41..c11eb93ed 100644 --- a/test/vector/cholesky.jl +++ b/test/vector/cholesky.jl @@ -3,8 +3,8 @@ module VBCholeskyTests using Distributions using LinearAlgebra using Test +using ADTypes: AutoEnzyme, AutoMooncake, AutoMooncakeForward, AutoReverseDiff using Bijectors.VectorBijectors -import DifferentiationInterface as DI using ForwardDiff: ForwardDiff using ReverseDiff: ReverseDiff using Mooncake: Mooncake @@ -13,17 +13,17 @@ using Enzyme: Enzyme, set_runtime_activity, Const, Forward, Reverse # Need runtime activity for some reason. # TODO(penelopeysm): Report upstream const adtypes = [ - DI.AutoReverseDiff(), - DI.AutoReverseDiff(; compile=true), - DI.AutoMooncake(), - DI.AutoMooncakeForward(), - DI.AutoEnzyme(; mode=set_runtime_activity(Forward), function_annotation=Const), - DI.AutoEnzyme(; mode=set_runtime_activity(Reverse), function_annotation=Const), + AutoReverseDiff(), + AutoReverseDiff(; compile=true), + AutoMooncake(), + AutoMooncakeForward(), + AutoEnzyme(; mode=set_runtime_activity(Forward), function_annotation=Const), + AutoEnzyme(; mode=set_runtime_activity(Reverse), function_annotation=Const), ] dists = [ - # Note: can't test LKJCholesky(1, ...) because its linked vector is length-zero and - # DifferentiationInterface trips up with empty vectors. + LKJCholesky(1, 1.0, 'U'), + LKJCholesky(1, 1.0, 'L'), LKJCholesky(3, 1.0, 'U'), LKJCholesky(3, 1.0, 'L'), LKJCholesky(5, 1.0, 'U'), diff --git a/test/vector/matrix.jl b/test/vector/matrix.jl index 53c5c8109..30d1f3c6e 100644 --- a/test/vector/matrix.jl +++ b/test/vector/matrix.jl @@ -4,7 +4,7 @@ using Distributions using LinearAlgebra using Test using PDMats -import DifferentiationInterface as DI +using ADTypes: AutoMooncake, AutoMooncakeForward using Bijectors.VectorBijectors using Enzyme: Enzyme using ForwardDiff: ForwardDiff @@ -19,7 +19,7 @@ M = [1 2 3; 4 5 6] # TODO(penelopeysm): ReverseDiff gives wrong results when differentiating # through VecCorrBijector. Correctness tests are disabled for now. # https://github.com/TuringLang/Bijectors.jl/issues/434 -lkj_test_adtypes = [DI.AutoMooncake(), DI.AutoMooncakeForward()] +lkj_test_adtypes = [AutoMooncake(), AutoMooncakeForward()] # Don't check that from_linked_vec(d)(randn(...)) is in support for LKJ, # The reason is because the inverse bijector for LKJ causes the diagonal diff --git a/test/vector/order.jl b/test/vector/order.jl index 682bb1f0d..9ae058e34 100644 --- a/test/vector/order.jl +++ b/test/vector/order.jl @@ -3,9 +3,9 @@ module VBOrderTests using Distributions using LinearAlgebra using Test +using ADTypes: AutoEnzyme, AutoMooncake, AutoMooncakeForward using Bijectors.VectorBijectors using Bijectors: ordered -import DifferentiationInterface as DI using Enzyme: Enzyme using ForwardDiff: ForwardDiff using ReverseDiff: ReverseDiff @@ -24,10 +24,10 @@ base_dists = [ # because of the heavy setindex! usage. # https://github.com/JuliaDiff/ReverseDiff.jl/issues/43 We just avoid testing it for now. joint_test_adtypes = [ - DI.AutoMooncake(), - DI.AutoMooncakeForward(), - DI.AutoEnzyme(; mode=Enzyme.Forward), - DI.AutoEnzyme(; mode=Enzyme.Reverse), + AutoMooncake(), + AutoMooncakeForward(), + AutoEnzyme(; mode=Enzyme.Forward), + AutoEnzyme(; mode=Enzyme.Reverse), ] @testset "Order statistics" begin diff --git a/test/vector/product.jl b/test/vector/product.jl index ecf5f680e..2137bcc76 100644 --- a/test/vector/product.jl +++ b/test/vector/product.jl @@ -4,28 +4,28 @@ using Distributions using LinearAlgebra using FillArrays: Fill using Test +using ADTypes: AutoEnzyme, AutoMooncake, AutoMooncakeForward, AutoReverseDiff using Bijectors.VectorBijectors -import DifferentiationInterface as DI using ForwardDiff: ForwardDiff using ReverseDiff: ReverseDiff using Mooncake: Mooncake using Enzyme: Enzyme, set_runtime_activity, Const, Forward, Reverse adtypes = [ - DI.AutoReverseDiff(), - DI.AutoReverseDiff(; compile=true), - DI.AutoMooncake(), - DI.AutoMooncakeForward(), + AutoReverseDiff(), + AutoReverseDiff(; compile=true), + AutoMooncake(), + AutoMooncakeForward(), # Need runtime activity for some reason. # TODO(penelopeysm): Report upstream - DI.AutoEnzyme(; mode=set_runtime_activity(Forward), function_annotation=Const), - DI.AutoEnzyme(; mode=set_runtime_activity(Reverse), function_annotation=Const), + AutoEnzyme(; mode=set_runtime_activity(Forward), function_annotation=Const), + AutoEnzyme(; mode=set_runtime_activity(Reverse), function_annotation=Const), ] # Enzyme segfaults on 1.12 + Windows. # https://github.com/EnzymeAD/Enzyme.jl/issues/2986 if VERSION >= v"1.12-" && Sys.iswindows() - filter!(a -> !(a isa DI.AutoEnzyme), adtypes) + filter!(a -> !(a isa AutoEnzyme), adtypes) end # These are purposely chosen because the vec_length output is the same but @@ -116,7 +116,7 @@ enzyme_failures = [ ) end - no_enzyme_adtypes = filter(adtype -> !(adtype isa DI.AutoEnzyme), adtypes) + no_enzyme_adtypes = filter(adtype -> !(adtype isa AutoEnzyme), adtypes) for d in enzyme_failures VectorBijectors.test_all( d; diff --git a/test/vector/reshaped.jl b/test/vector/reshaped.jl index 64d07ce99..9674d2fce 100644 --- a/test/vector/reshaped.jl +++ b/test/vector/reshaped.jl @@ -3,8 +3,8 @@ module VBReshapedTests using Distributions using LinearAlgebra using Test +using ADTypes: AutoEnzyme, AutoMooncake, AutoMooncakeForward, AutoReverseDiff using Bijectors.VectorBijectors -import DifferentiationInterface as DI using Enzyme: Enzyme using ForwardDiff: ForwardDiff using ReverseDiff: ReverseDiff @@ -30,11 +30,11 @@ reshaped = [ # Fails on 1.10: https://github.com/EnzymeAD/Enzyme.jl/issues/2987 adtypes_no_enz_rvs = [ - DI.AutoReverseDiff(), - DI.AutoReverseDiff(; compile=true), - DI.AutoMooncake(), - DI.AutoMooncakeForward(), - DI.AutoEnzyme(; mode=Enzyme.Forward, function_annotation=Enzyme.Const), + AutoReverseDiff(), + AutoReverseDiff(; compile=true), + AutoMooncake(), + AutoMooncakeForward(), + AutoEnzyme(; mode=Enzyme.Forward, function_annotation=Enzyme.Const), ] reshaped_no_enzyme = [reshape(Beta(2, 2), (1, 1, 1, 1, 1))]