Merge remote-tracking branch 'origin/main' into gd/forward_mooncake

Author: gdalle

.buildkite/pipeline.yml

@@ -1,6 +1,5 @@
 steps:
   - label: "DI GPU tests"
-    if: build.pull_request.labels includes "gpu"
     plugins:
       - JuliaCI/julia#v1:
           version: "1"

.github/workflows/Test.yml

@@ -20,7 +20,7 @@ jobs:
     name: ${{ matrix.version }} - DI (${{ matrix.group }})
     runs-on: ubuntu-latest
     if: ${{ !contains(github.event.pull_request.labels.*.name, 'skipci') }}
-    timeout-minutes: 60
+    timeout-minutes: 120
     permissions: # needed to allow julia-actions/cache to proactively delete old caches that it has created
       actions: write
       contents: read

DifferentiationInterface/CHANGELOG.md

@@ -7,6 +7,32 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [0.7.4]
+
+### Added
+
+- Make `AutoForwardFromPrimitive` and `AutoReverseFromPrimitive` public ([#825])
+
+### Fixed
+
+- Replace `one` with `oneunit` in basis computation ([#826])
+
+## [0.7.3]
+
+### Fixed
+
+- Bump compat for SparseConnectivityTracer v1 ([#823])
+
+## [0.7.2]
+
+### Feat
+
+- Backend switching for Mooncake ([#768])
+
+### Fixed
+
+- Speed up sparse preparation for GPU arrays ([#818])
+
 ## [0.7.1]
 
 ### Feat
@@ -38,17 +64,25 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Allocate Enzyme shadow memory during preparation ([#782])
 
-[unreleased]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.7.1...main
+[unreleased]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.7.4...main
+[0.7.4]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.7.3...DifferentiationInterface-v0.7.4
+[0.7.3]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.7.2...DifferentiationInterface-v0.7.3
+[0.7.2]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.7.1...DifferentiationInterface-v0.7.2
 [0.7.1]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.7.0...DifferentiationInterface-v0.7.1
 [0.7.0]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.6.54...DifferentiationInterface-v0.7.0
 [0.6.54]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.6.53...DifferentiationInterface-v0.6.54
 [0.6.53]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.6.52...DifferentiationInterface-v0.6.53
 
+[#826]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/826
+[#825]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/825
+[#823]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/823
+[#818]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/818
 [#812]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/812
 [#810]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/810
 [#809]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/809
 [#799]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/799
 [#795]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/795
 [#790]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/790
 [#788]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/788
-[#782]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/782
+[#782]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/782
+[#768]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/768

DifferentiationInterface/Project.toml

@@ -1,7 +1,7 @@
 name = "DifferentiationInterface"
 uuid = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 authors = ["Guillaume Dalle", "Adrian Hill"]
-version = "0.7.1"
+version = "0.7.4"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -42,7 +42,9 @@ DifferentiationInterfaceGPUArraysCoreExt = "GPUArraysCore"
 DifferentiationInterfaceGTPSAExt = "GTPSA"
 DifferentiationInterfaceMooncakeExt = "Mooncake"
 DifferentiationInterfacePolyesterForwardDiffExt = [
-    "PolyesterForwardDiff", "ForwardDiff", "DiffResults"
+    "PolyesterForwardDiff",
+    "ForwardDiff",
+    "DiffResults",
 ]
 DifferentiationInterfaceReverseDiffExt = ["ReverseDiff", "DiffResults"]
 DifferentiationInterfaceSparseArraysExt = "SparseArrays"
@@ -54,11 +56,12 @@ DifferentiationInterfaceTrackerExt = "Tracker"
 DifferentiationInterfaceZygoteExt = ["Zygote", "ForwardDiff"]
 
 [compat]
-Aqua = "0.8.12"
 ADTypes = "1.13.0"
+Aqua = "0.8.12"
 ChainRulesCore = "1.23.0"
 ComponentArrays = "0.15.27"
 DataFrames = "1.7.0"
+Dates = "1"
 DiffResults = "1.1.0"
 Diffractor = "=0.2.6"
 Enzyme = "0.13.39"
@@ -80,7 +83,7 @@ PolyesterForwardDiff = "0.1.2"
 Random = "1"
 ReverseDiff = "1.15.1"
 SparseArrays = "1"
-SparseConnectivityTracer = "0.6.14"
+SparseConnectivityTracer = "0.6.14, 1"
 SparseMatrixColorings = "0.4.9"
 StableRNGs = "1.0.1"
 StaticArrays = "1.9.7"
@@ -96,6 +99,7 @@ Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
+Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 Diffractor = "9f5e2b26-1114-432f-b630-d3fe2085c51c"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
@@ -128,6 +132,7 @@ test = [
     "Aqua",
     "ComponentArrays",
     "DataFrames",
+    "Dates",
     "ExplicitImports",
     "JET",
     "JLArrays",

DifferentiationInterface/docs/src/api.md

@@ -132,6 +132,13 @@ MixedMode
 DenseSparsityDetector
 ```
 
+### From primitive
+
+```@docs
+DifferentiationInterface.AutoForwardFromPrimitive
+DifferentiationInterface.AutoReverseFromPrimitive
+```
+
 ## Internals
 
 The following is not part of the public API.

DifferentiationInterface/docs/src/explanation/advanced.md

@@ -86,7 +86,16 @@ AutoSparse(
 )
 ```
 
-At the moment, mixed mode tends to work best when the [`GreedyColoringAlgorithm`](@extref SparseMatrixColorings.GreedyColoringAlgorithm) is provided with a [`RandomOrder`](@extref SparseMatrixColorings.RandomOrder) instead of the usual [`NaturalOrder`](@extref SparseMatrixColorings.NaturalOrder).
+At the moment, mixed mode tends to work best (output fewer colors) when the [`GreedyColoringAlgorithm`](@extref SparseMatrixColorings.GreedyColoringAlgorithm) is provided with a [`RandomOrder`](@extref SparseMatrixColorings.RandomOrder) instead of the usual [`NaturalOrder`](@extref SparseMatrixColorings.NaturalOrder), and when "post-processing" is activated after coloring.
+For full reproducibility, you should use a random number generator from [StableRNGs.jl](https://github.com/JuliaRandom/StableRNGs.jl).
+Thus, the right setup looks like:
+
+```julia
+using StableRNGs
+
+seed = 3
+coloring_algorithm = GreedyColoringAlgorithm(RandomOrder(StableRNG(seed), seed); postprocessing=true)
+```
 
 ## Batch mode
 

DifferentiationInterface/docs/src/explanation/operators.md

@@ -152,4 +152,4 @@ For same-point preparation, the same rules hold with two modifications:
 
 !!! warning
     These rules hold for the majority of backends, but there are some exceptions.
-    The most important exception is [ReverseDiff](@ref) and its taping mechanism, which is sensitive to control flow inside the function.
+    The most important exception is [ReverseDiff.jl](https://github.com/JuliaDiff/ReverseDiff.jl) and its taping mechanism, which is sensitive to control flow inside the function.

DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/onearg.jl

@@ -189,27 +189,27 @@ end
 function DI.value_and_derivative(
     f::F, backend::AutoForwardDiff, x, contexts::Vararg{DI.Context,C}
 ) where {F,C}
-    y, ty = DI.value_and_pushforward(f, backend, x, (one(x),), contexts...)
+    y, ty = DI.value_and_pushforward(f, backend, x, (oneunit(x),), contexts...)
     return y, only(ty)
 end
 
 function DI.value_and_derivative!(
     f::F, der, backend::AutoForwardDiff, x, contexts::Vararg{DI.Context,C}
 ) where {F,C}
-    y, _ = DI.value_and_pushforward!(f, (der,), backend, x, (one(x),), contexts...)
+    y, _ = DI.value_and_pushforward!(f, (der,), backend, x, (oneunit(x),), contexts...)
     return y, der
 end
 
 function DI.derivative(
     f::F, backend::AutoForwardDiff, x, contexts::Vararg{DI.Context,C}
 ) where {F,C}
-    return only(DI.pushforward(f, backend, x, (one(x),), contexts...))
+    return only(DI.pushforward(f, backend, x, (oneunit(x),), contexts...))
 end
 
 function DI.derivative!(
     f::F, der, backend::AutoForwardDiff, x, contexts::Vararg{DI.Context,C}
 ) where {F,C}
-    DI.pushforward!(f, (der,), backend, x, (one(x),), contexts...)
+    DI.pushforward!(f, (der,), backend, x, (oneunit(x),), contexts...)
     return der
 end
 
@@ -220,7 +220,7 @@ function DI.prepare_derivative_nokwarg(
 ) where {F,C}
     _sig = DI.signature(f, backend, x, contexts...; strict)
     pushforward_prep = DI.prepare_pushforward_nokwarg(
-        strict, f, backend, x, (one(x),), contexts...
+        strict, f, backend, x, (oneunit(x),), contexts...
     )
     return ForwardDiffOneArgDerivativePrep(_sig, pushforward_prep)
 end
@@ -234,7 +234,7 @@ function DI.value_and_derivative(
 ) where {F,C}
     DI.check_prep(f, prep, backend, x, contexts...)
     y, ty = DI.value_and_pushforward(
-        f, prep.pushforward_prep, backend, x, (one(x),), contexts...
+        f, prep.pushforward_prep, backend, x, (oneunit(x),), contexts...
     )
     return y, only(ty)
 end
@@ -249,7 +249,7 @@ function DI.value_and_derivative!(
 ) where {F,C}
     DI.check_prep(f, prep, backend, x, contexts...)
     y, _ = DI.value_and_pushforward!(
-        f, (der,), prep.pushforward_prep, backend, x, (one(x),), contexts...
+        f, (der,), prep.pushforward_prep, backend, x, (oneunit(x),), contexts...
     )
     return y, der
 end
@@ -263,7 +263,7 @@ function DI.derivative(
 ) where {F,C}
     DI.check_prep(f, prep, backend, x, contexts...)
     return only(
-        DI.pushforward(f, prep.pushforward_prep, backend, x, (one(x),), contexts...)
+        DI.pushforward(f, prep.pushforward_prep, backend, x, (oneunit(x),), contexts...)
     )
 end
 
@@ -276,7 +276,9 @@ function DI.derivative!(
     contexts::Vararg{DI.Context,C},
 ) where {F,C}
     DI.check_prep(f, prep, backend, x, contexts...)
-    DI.pushforward!(f, (der,), prep.pushforward_prep, backend, x, (one(x),), contexts...)
+    DI.pushforward!(
+        f, (der,), prep.pushforward_prep, backend, x, (oneunit(x),), contexts...
+    )
     return der
 end
 
@@ -638,9 +640,9 @@ function DI.second_derivative(
 ) where {F,C}
     DI.check_prep(f, prep, backend, x, contexts...)
     T = tag_type(f, backend, x)
-    xdual = make_dual(T, x, one(x))
+    xdual = make_dual(T, x, oneunit(x))
     T2 = tag_type(f, backend, xdual)
-    xdual2 = make_dual(T2, xdual, one(xdual))
+    xdual2 = make_dual(T2, xdual, oneunit(xdual))
     contexts_dual = translate(typeof(xdual2), contexts)
     ydual = f(xdual2, contexts_dual...)
     return myderivative(T, myderivative(T2, ydual))
@@ -656,9 +658,9 @@ function DI.second_derivative!(
 ) where {F,C}
     DI.check_prep(f, prep, backend, x, contexts...)
     T = tag_type(f, backend, x)
-    xdual = make_dual(T, x, one(x))
+    xdual = make_dual(T, x, oneunit(x))
     T2 = tag_type(f, backend, xdual)
-    xdual2 = make_dual(T2, xdual, one(xdual))
+    xdual2 = make_dual(T2, xdual, oneunit(xdual))
     contexts_dual = translate(typeof(xdual2), contexts)
     ydual = f(xdual2, contexts_dual...)
     return myderivative!(T, der2, myderivative(T2, ydual))
@@ -673,9 +675,9 @@ function DI.value_derivative_and_second_derivative(
 ) where {F,C}
     DI.check_prep(f, prep, backend, x, contexts...)
     T = tag_type(f, backend, x)
-    xdual = make_dual(T, x, one(x))
+    xdual = make_dual(T, x, oneunit(x))
     T2 = tag_type(f, backend, xdual)
-    xdual2 = make_dual(T2, xdual, one(xdual))
+    xdual2 = make_dual(T2, xdual, oneunit(xdual))
     contexts_dual = translate(typeof(xdual2), contexts)
     ydual = f(xdual2, contexts_dual...)
     y = myvalue(T, myvalue(T2, ydual))
@@ -695,9 +697,9 @@ function DI.value_derivative_and_second_derivative!(
 ) where {F,C}
     DI.check_prep(f, prep, backend, x, contexts...)
     T = tag_type(f, backend, x)
-    xdual = make_dual(T, x, one(x))
+    xdual = make_dual(T, x, oneunit(x))
     T2 = tag_type(f, backend, xdual)
-    xdual2 = make_dual(T2, xdual, one(xdual))
+    xdual2 = make_dual(T2, xdual, oneunit(xdual))
     contexts_dual = translate(typeof(xdual2), contexts)
     ydual = f(xdual2, contexts_dual...)
     y = myvalue(T, myvalue(T2, ydual))
@@ -756,7 +758,7 @@ function DI.value_gradient_and_hessian!(
         contexts isa NTuple{C,DI.GeneralizedConstant}
     )
         fc = DI.fix_tail(f, map(DI.unwrap, contexts)...)
-        result = DiffResult(one(eltype(x)), (grad, hess))
+        result = DiffResult(oneunit(eltype(x)), (grad, hess))
         result = hessian!(result, fc, x)
         y = DR.value(result)
         grad === DR.gradient(result) || copyto!(grad, DR.gradient(result))
@@ -855,7 +857,7 @@ function DI.value_gradient_and_hessian!(
     DI.check_prep(f, prep, backend, x, contexts...)
     contexts_dual = translate_prepared(contexts, prep.contexts_dual)
     fc = DI.fix_tail(f, contexts_dual...)
-    result = DiffResult(one(eltype(x)), (grad, hess))
+    result = DiffResult(oneunit(eltype(x)), (grad, hess))
     CHK = tag_type(backend) === Nothing
     if CHK
         checktag(prep.result_config, f, x)

DifferentiationInterface/ext/DifferentiationInterfaceGPUArraysCoreExt/DifferentiationInterfaceGPUArraysCoreExt.jl

@@ -6,16 +6,14 @@ using GPUArraysCore: @allowscalar, AbstractGPUArray
 function DI.basis(a::AbstractGPUArray{T}, i) where {T}
     b = similar(a)
     fill!(b, zero(T))
-    @allowscalar b[i] = one(T)
+    @allowscalar b[i] = oneunit(T)
     return b
 end
 
 function DI.multibasis(a::AbstractGPUArray{T}, inds) where {T}
     b = similar(a)
     fill!(b, zero(T))
-    for i in inds
-        @allowscalar b[i] = one(T)
-    end
+    view(b, inds) .= oneunit(T)
     return b
 end
 

DifferentiationInterface/src/DifferentiationInterface.jl

@@ -128,6 +128,7 @@ export AutoSparse
 ## Public but not exported
 
 @public inner, outer
+@public AutoForwardFromPrimitive, AutoReverseFromPrimitive
 
 include("init.jl")