changes from reviews - 2

Author: AstitvaAggarwal

DifferentiationInterface/docs/src/explanation/backends.md

@@ -95,7 +95,7 @@ In general, using a forward outer backend over a reverse inner backend will yiel
 The wrapper [`DifferentiateWith`](@ref) allows you to switch between backends.
 It takes a function `f` and specifies that `f` should be differentiated with the substitute backend of your choice, instead of whatever true backend the surrounding code is trying to use.
 In other words, when someone tries to differentiate `dw = DifferentiateWith(f, substitute_backend)` with `true_backend`, then `substitute_backend` steps in and `true_backend` does not dive into the function `f` itself.
-At the moment, `DifferentiateWith` only works when `true_backend` is either [ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl), [ChainRules.jl](https://github.com/JuliaDiff/ChainRules.jl) or a [Mooncake.jl](https://github.com/compintell/Mooncake.jl)-compatible backend.
+At the moment, `DifferentiateWith` only works when `true_backend` is either [ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl), [Mooncake.jl](https://github.com/compintell/Mooncake.jl), or a [ChainRules.jl](https://github.com/JuliaDiff/ChainRules.jl)-compatible backend (e.g., [Zygote.jl](https://github.com/FluxML/Zygote.jl)).
 
 ## Implementations
 

DifferentiationInterface/docs/src/faq/differentiability.md

@@ -111,4 +111,4 @@ There are, however, translation utilities:
 ### Backend switch
 
 Also note the existence of [`DifferentiationInterface.DifferentiateWith`](@ref), which allows the user to wrap a function that should be differentiated with a specific backend.
-Right now it only targets ForwardDiff.jl, ChainRulesCore.jl and Mooncake.jl but PRs are welcome to define Enzyme.jl rules for this object.
+Right now, it only targets ForwardDiff.jl, Mooncake.jl, ChainRules.jl-compatible backends (e.g., [Zygote.jl](https://github.com/FluxML/Zygote.jl)),  but PRs are welcome to define Enzyme.jl rules for this object.

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/differentiate_with.jl

@@ -7,18 +7,18 @@ function Mooncake.rrule!!(dw::CoDual{<:DI.DifferentiateWith}, x::CoDual{<:Number
     y = zero_fcodual(f(primal_x))
 
     # output is a vector, so we need to use the vector pullback
-    function pullback!!(dy::NoRData)
+    function pullback_array!!(dy::NoRData)
         tx = DI.pullback(f, backend, primal_x, (fdata(y.dx),))
         return NoRData(), only(tx)
     end
 
     # output is a scalar, so we can use the scalar pullback
-    function pullback!!(dy::Number)
+    function pullback_scalar!!(dy::Number)
         tx = DI.pullback(f, backend, primal_x, (dy,))
         return NoRData(), only(tx)
     end
 
-    return y, pullback!!
+    return y, typeof(primal(y)) <: Number ? pullback_scalar!! : pullback_array!!
 end
 
 function Mooncake.rrule!!(dw::CoDual{<:DI.DifferentiateWith}, x::CoDual{<:AbstractArray})
@@ -27,22 +27,20 @@ function Mooncake.rrule!!(dw::CoDual{<:DI.DifferentiateWith}, x::CoDual{<:Abstra
     fdata_arg = fdata(x.dx)
     (; f, backend) = primal_func
     y = zero_fcodual(f(primal_x))
-    # in case x is mutated in f calls
-    cp_primal_x = copy(primal_x)
 
     # output is a vector, so we need to use the vector pullback
-    function pullback!!(dy::NoRData)
-        tx = DI.pullback(f, backend, cp_primal_x, (fdata(y.dx),))
+    function pullback_array!!(dy::NoRData)
+        tx = DI.pullback(f, backend, primal_x, (fdata(y.dx),))
         fdata_arg .+= only(tx)
         return NoRData(), dy
     end
 
     # output is a scalar, so we can use the scalar pullback
-    function pullback!!(dy::Number)
-        tx = DI.pullback(f, backend, cp_primal_x, (dy,))
+    function pullback_scalar!!(dy::Number)
+        tx = DI.pullback(f, backend, primal_x, (dy,))
         fdata_arg .+= only(tx)
         return NoRData(), NoRData()
     end
 
-    return y, pullback!!
+    return y, typeof(primal(y)) <: Number ? pullback_scalar!! : pullback_array!!
 end