feat: make FromPrimitive wrappers public

Author: gdalle

DifferentiationInterface/CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- Make `AutoForwardFromPrimitive` and `AutoReverseFromPrimitive` public ([#824])
+
 ## [0.7.3]
 
 ### Fixed
@@ -62,6 +66,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.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
 
+[#824]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/824
 [#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

DifferentiationInterface/docs/src/api.md

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

DifferentiationInterface/src/DifferentiationInterface.jl

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

DifferentiationInterface/src/misc/from_primitive.jl

@@ -12,17 +12,21 @@ function pick_batchsize(backend::FromPrimitive, N::Integer)
 end
 
 """
-    AutoForwardFromPrimitive
+    AutoForwardFromPrimitive(backend::AbstractADType)
 
-Wrapper which forces a given backend to act as a reverse-mode backend.
+Wrapper which forces a given backend to act as a forward-mode backend, using only its native `value_and_pushforward` primitive and re-implementing the rest from scratch.
 
-Used in internal testing.
+!!! tip
+    This can be useful to circumvent high-level operators when they have impractical limitations.
+    For instance, ForwardDiff.jl's `jacobian` does not support GPU arrays but its `pushforward` does, so `AutoForwardFromPrimitive(AutoForwardDiff())` has a GPU-friendly `jacobian`.
 """
 struct AutoForwardFromPrimitive{inplace,B<:AbstractADType} <: FromPrimitive{inplace}
     backend::B
 end
 
-function AutoForwardFromPrimitive(backend::AbstractADType; inplace=true)
+function AutoForwardFromPrimitive(
+    backend::AbstractADType; inplace::Bool=Bool(inplace_support(backend))
+)
     return AutoForwardFromPrimitive{inplace,typeof(backend)}(backend)
 end
 
@@ -133,17 +137,17 @@ function value_and_pushforward!(
 end
 
 """
-    AutoReverseFromPrimitive
+    AutoReverseFromPrimitive(backend::AbstractADType)
 
-Wrapper which forces a given backend to act as a reverse-mode backend.
-
-Used in internal testing.
+Wrapper which forces a given backend to act as a reverse-mode backend, using only its native `value_and_pullback` implementation and rebuilding the rest from scratch.
 """
 struct AutoReverseFromPrimitive{inplace,B<:AbstractADType} <: FromPrimitive{inplace}
     backend::B
 end
 
-function AutoReverseFromPrimitive(backend::AbstractADType; inplace=true)
+function AutoReverseFromPrimitive(
+    backend::AbstractADType; inplace::Bool=Bool(inplace_support(backend))
+)
     return AutoReverseFromPrimitive{inplace,typeof(backend)}(backend)
 end
 

DifferentiationInterface/test/Core/SimpleFiniteDiff/test.jl

@@ -141,6 +141,12 @@ end
 
 @testset "Weird arrays" begin
     test_differentiation(
-        AutoSimpleFiniteDiff(), vcat(static_scenarios(), gpu_scenarios()); logging=LOGGING
+        [
+            AutoSimpleFiniteDiff(),
+            AutoForwardFromPrimitive(AutoSimpleFiniteDiff()),
+            AutoReverseFromPrimitive(AutoSimpleFiniteDiff()),
+        ],
+        vcat(static_scenarios(), gpu_scenarios());
+        logging=LOGGING,
     )
 end;