Bump deps to latest, route AD through AbstractPPL instead of DifferentiationInterface

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.

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"

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"

docs/src/usage.md

@@ -37,4 +37,8 @@ flow_trained, stats, _ = train_flow(
     ADbackend=ADTypes.AutoMooncake(; config=Mooncake.Config()),
     show_progress=true,
 )
-```
+```
+
+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.

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"

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")

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

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

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"

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]

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"