diff --git a/HISTORY.md b/HISTORY.md new file mode 100644 index 00000000..f25540ca --- /dev/null +++ b/HISTORY.md @@ -0,0 +1,15 @@ +# 0.3.0 + +## Breaking changes + +DifferentiationInterface has been removed as a dependency. Automatic differentiation is now routed through AbstractPPL's evaluator interface (`AbstractPPL.prepare` and `AbstractPPL.value_and_gradient!!`), following the rest of the Turing ecosystem. + +The AD backend package must now be loaded so that its `AbstractPPL.prepare` method is available: + + - `AutoForwardDiff` works with `using ForwardDiff`. + - `AutoMooncake` works with `using Mooncake`. + - Other backends routed through DifferentiationInterface (`AutoZygote`, `AutoReverseDiff`, `AutoEnzyme`) additionally require `using DifferentiationInterface` alongside the concrete backend package. + +## Other changes + +Bijectors compat now includes 0.16 and CUDA compat now includes 6. diff --git a/Project.toml b/Project.toml index 109b2122..29e2553a 100644 --- a/Project.toml +++ b/Project.toml @@ -1,11 +1,11 @@ name = "NormalizingFlows" uuid = "50e4474d-9f12-44b7-af7a-91ab30ff6256" -version = "0.2.2" +version = "0.3.0" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" +AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" Bijectors = "76274a88-744f-5084-9051-94815aaf08c4" -DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c" @@ -25,9 +25,9 @@ NormalizingFlowsCUDAExt = "CUDA" [compat] ADTypes = "1" +AbstractPPL = "0.15.4" Bijectors = "0.12.6, 0.13, 0.14, 0.15, 0.16" CUDA = "5, 6.2" -DifferentiationInterface = "0.6, 0.7" Distributions = "0.25" DocStringExtensions = "0.9" Flux = "0.16" diff --git a/docs/Project.toml b/docs/Project.toml index 900b5371..1241a55b 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -16,4 +16,4 @@ Documenter = "1.17.0" Flux = "0.16.10" Functors = "0.5.2" LiveServer = "1.5.0" -NormalizingFlows = "0.2.2" +NormalizingFlows = "0.3.0" diff --git a/docs/src/usage.md b/docs/src/usage.md index cefb8167..9e62f994 100644 --- a/docs/src/usage.md +++ b/docs/src/usage.md @@ -37,4 +37,8 @@ flow_trained, stats, _ = train_flow( ADbackend=ADTypes.AutoMooncake(; config=Mooncake.Config()), show_progress=true, ) -``` \ No newline at end of file +``` + +Gradients are computed through AbstractPPL's evaluator interface, so the chosen AD backend package must be loaded first. +`AutoForwardDiff` needs `using ForwardDiff`, and `AutoMooncake` needs `using Mooncake` (as above). +The other backends (`AutoZygote`, `AutoReverseDiff`, `AutoEnzyme`) additionally need `using DifferentiationInterface` alongside the backend package. \ No newline at end of file diff --git a/example/Project.toml b/example/Project.toml index 8b44b7b3..dd3a6981 100644 --- a/example/Project.toml +++ b/example/Project.toml @@ -34,7 +34,7 @@ Functors = "0.5.2" IrrationalConstants = "0.2.6" LogDensityProblems = "2.2.0" Mooncake = "0.4.142, 0.5" -NormalizingFlows = "0.2.2" +NormalizingFlows = "0.3.0" Optimisers = "0.4.7" Plots = "1.41.6" ProgressMeter = "1.11.0" diff --git a/example/demo_neural_spline_flow.jl b/example/demo_neural_spline_flow.jl index 4e98c308..29d7a894 100644 --- a/example/demo_neural_spline_flow.jl +++ b/example/demo_neural_spline_flow.jl @@ -1,7 +1,7 @@ using Random, Distributions, LinearAlgebra using Functors using Optimisers, ADTypes -using Zygote +using Zygote, DifferentiationInterface using NormalizingFlows include("SyntheticTargets.jl") diff --git a/src/NormalizingFlows.jl b/src/NormalizingFlows.jl index 70c46996..2f754377 100644 --- a/src/NormalizingFlows.jl +++ b/src/NormalizingFlows.jl @@ -10,7 +10,7 @@ using StatsBase using Bijectors using Bijectors: PartitionMask, Inverse, combine, partition using Functors -import DifferentiationInterface as DI +using AbstractPPL: AbstractPPL using DocStringExtensions @@ -31,15 +31,18 @@ Arguments # Keyword Arguments - `max_iters::Int=1000`: maximum number of iterations - `optimiser::Optimisers.AbstractRule=Optimisers.ADAM()`: optimiser to compute the steps -- `ADbackend::ADTypes.AbstractADType=ADTypes.AutoZygote()`: +- `ADbackend::ADTypes.AbstractADType=ADTypes.AutoZygote()`: automatic differentiation backend, currently supports - `ADTypes.AutoZygote()`, `ADTypes.ForwardDiff()`, `ADTypes.ReverseDiff()`, + `ADTypes.AutoZygote()`, `ADTypes.ForwardDiff()`, `ADTypes.ReverseDiff()`, `ADTypes.AutoMooncake()` and `ADTypes.AutoEnzyme(; mode=Enzyme.set_runtime_activity(Enzyme.Reverse), function_annotation=Enzyme.Const, )`. If user wants to use `AutoEnzyme`, please make sure to include the `set_runtime_activity` and `function_annotation` as shown above. + Gradients are computed through AbstractPPL's evaluator interface, so the chosen backend package must be loaded first. + `AutoForwardDiff` needs `using ForwardDiff`, and `AutoMooncake` needs `using Mooncake`. + The other backends (`AutoZygote`, `AutoReverseDiff`, `AutoEnzyme`) additionally need `using DifferentiationInterface` alongside the backend package. - `kwargs...`: additional keyword arguments for `optimize` (See [`optimize`](@ref) for details) # Returns diff --git a/src/optimize.jl b/src/optimize.jl index b4adad91..cc8dc82f 100644 --- a/src/optimize.jl +++ b/src/optimize.jl @@ -6,11 +6,14 @@ function pm_next!(pm, stats::NamedTuple) end function _prepare_gradient(loss, adbackend, θ, args...) - return DI.prepare_gradient(loss, adbackend, θ, map(DI.Constant, args)...) + return AbstractPPL.prepare(adbackend, loss, θ; context=args) end function _value_and_gradient(loss, prep, adbackend, θ, args...) - return DI.value_and_gradient(loss, prep, adbackend, θ, map(DI.Constant, args)...) + val, grad = AbstractPPL.value_and_gradient!!(prep, θ) + # value_and_gradient!! may return a gradient aliasing the prep's internal + # buffer, which the next call overwrites + return val, copy(grad) end """ @@ -31,7 +34,7 @@ Iteratively updating the parameters `θ` of the normalizing flow `re(θ)` by cal - `loss`: a general loss function θ -> loss(θ, args...) returning a scalar loss value that will be minimised - `θ₀::AbstractVector{T}`: initial parameters for the loss function (in the context of normalizing flows, it will be the flattened flow parameters) - `re`: reconstruction function that maps the flattened parameters to the normalizing flow -- `args...`: additional arguments for `loss` (will be set as DI.Constant) +- `args...`: additional arguments for `loss` (treated as constants during differentiation) # Keyword Arguments - `max_iters::Int=10000`: maximum number of iterations diff --git a/test/Project.toml b/test/Project.toml index b45ccaa8..8b04e703 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -22,13 +22,13 @@ ADTypes = "1.22.2" Bijectors = "0.16.2" DifferentiationInterface = "0.7.20" Distributions = "0.25.129" -Enzyme = "0.13.186" +Enzyme = "0.13.186 - 0.13.188" Flux = "0.16.10" ForwardDiff = "1.4.1" Functors = "0.5.2" MonotonicSplines = "0.3.3" Mooncake = "0.4.142, 0.5" -NormalizingFlows = "0.2.2" +NormalizingFlows = "0.3.0" Optimisers = "0.4.7" ReverseDiff = "1.17.0" Zygote = "0.7.11" diff --git a/test/ad.jl b/test/ad.jl index b8a42117..2adf5966 100644 --- a/test/ad.jl +++ b/test/ad.jl @@ -1,4 +1,4 @@ -@testset "DI.AD with context wrapper" begin +@testset "AD with context wrapper" begin f(x, y, z) = sum(abs2, x .+ y .+ z) @testset "$T" for T in [Float32, Float64] diff --git a/test/ext/CUDA/Project.toml b/test/ext/CUDA/Project.toml index 15797947..58dc0bb2 100644 --- a/test/ext/CUDA/Project.toml +++ b/test/ext/CUDA/Project.toml @@ -12,4 +12,4 @@ Bijectors = "0.16.2" CUDA = "6.2.1" Distributions = "0.25.129" Flux = "0.16.10" -NormalizingFlows = "0.2.2" +NormalizingFlows = "0.3.0"