Use Enzyme.ignore_derivatives, now that landed (#1707)

* Use Enzyme.ignore_derivatives, now that landed

* move

* fix

* Apply suggestion from @github-actions[bot]

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Update ReactantABI to use EnzymeCore directly

* Change ignore_derivatives reference to EnzymeCore

* fix: docs

---------

Co-authored-by: Avik Pal <avikpal@mit.edu>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Author: 3 people

Project.toml

@@ -81,8 +81,8 @@ CUDA = "5.6"
 DLFP8Types = "0.1"
 Downloads = "1.6"
 EnumX = "1"
-Enzyme = "0.13.78"
-EnzymeCore = "0.8.13"
+Enzyme = "0.13.81"
+EnzymeCore = "0.8.14"
 FillArrays = "1.13"
 Float8s = "0.1"
 Functors = "0.5"

docs/Project.toml

@@ -2,6 +2,7 @@
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 DocumenterVitepress = "4710194d-e776-4893-9690-8d956a29c365"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 OpenSSL_jll = "458c3c95-2e84-50aa-8efc-19380b2a3a95"
 Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 ReactantCore = "a3311ec8-5e00-46d5-b541-4f83e724a433"

docs/make.jl

@@ -1,4 +1,4 @@
-using Reactant, ReactantCore
+using Reactant, ReactantCore, EnzymeCore
 using Documenter, DocumenterVitepress
 
 DocMeta.setdocmeta!(Reactant, :DocTestSetup, :(using Reactant); recursive=true)
@@ -7,6 +7,7 @@ makedocs(;
     modules=[
         Reactant,
         ReactantCore,
+        EnzymeCore,
         Reactant.XLA,
         Reactant.MLIR,
         Reactant.MLIR.API,

docs/src/api/api.md

@@ -78,7 +78,7 @@ Reactant.addressable_devices
 ## Differentiation Specific API
 
 ```@docs
-Reactant.ignore_derivatives
+EnzymeCore.ignore_derivatives
 ```
 
 ## Persistent Compilation Cache

docs/src/api/internal.md

@@ -14,3 +14,14 @@ Reactant.Compiler.codegen_unflatten!
 Reactant.Compiler.codegen_flatten!
 Reactant.Compiler.codegen_xla_call
 ```
+
+## Other Docstrings
+
+!!! warning "Private"
+
+    These docstrings are present here to prevent missing docstring warnings. For official
+    Enzyme documentation checkout https://enzymead.github.io/Enzyme.jl/stable/.
+
+```@autodocs
+Modules = [EnzymeCore, EnzymeCore.EnzymeRules]
+```

docs/src/tutorials/automatic-differentiation.md

@@ -190,12 +190,12 @@ nothing # hide
 
 ### Ignoring Derivatives
 
-Use [`Reactant.ignore_derivatives`](@ref) to exclude parts of computation from gradient:
+Use [`EnzymeCore.ignore_derivatives`](@ref) to exclude parts of computation from gradient:
 
 ```@example autodiff_tutorial
 function func_with_ignore(x)
     # This part won't contribute to gradient
-    ignored_sum = Reactant.ignore_derivatives(sum(x))
+    ignored_sum = Enzyme.ignore_derivatives(sum(x))
     # This part will contribute
     return sum(x .^ 2) + ignored_sum
 end

src/Enzyme.jl

@@ -585,24 +585,4 @@ function overload_autodiff(
     end
 end
 
-"""
-    ignore_derivatives(args...)
-
-Prevents the flow of gradients (and higher-order derivatives) by creating a new value that
-is detached from the original value. This is an identity operation on the primal. This can
-be applied on a nested structure of arrays and we will apply the operation on each of the
-leaves.
-"""
-function ignore_derivatives(args...)
-    res = map(ignore_derivatives_internal, args)
-    length(args) == 1 && return only(res)
-    return res
-end
-
-function ignore_derivatives_internal(arg)
-    return Functors.fmap(arg) do argᵢ
-        argᵢ isa AnyTracedRArray && (argᵢ = materialize_traced_array(argᵢ))
-        argᵢ isa TracedType && return @opcall ignore_derivatives(argᵢ)
-        return argᵢ
-    end
-end
+const ignore_derivatives = EnzymeCore.ignore_derivatives

src/Overlay.jl

@@ -21,6 +21,19 @@ end
     return overload_autodiff(rmode, f, rt, args...)
 end
 
+@reactant_overlay function EnzymeCore.ignore_derivatives(args...)
+    res = map(args) do arg
+        return Functors.fmap(arg) do argᵢ
+            argᵢ isa AnyTracedRArray &&
+                (argᵢ = call_with_reactant(materialize_traced_array, argᵢ))
+            argᵢ isa TracedType && return @opcall ignore_derivatives(argᵢ)
+            return argᵢ
+        end
+    end
+    length(args) == 1 && return only(res)
+    return res
+end
+
 # Random.jl overlays
 @reactant_overlay @noinline function Random.default_rng()
     return call_with_reactant(TracedRandom.default_rng)

src/Reactant.jl

@@ -26,6 +26,7 @@ using Enzyme:
     DuplicatedNoNeed,
     EnzymeRules,
     Reverse
+using EnzymeCore: EnzymeCore
 
 export allowscalar, @allowscalar # re-exported from GPUArraysCore
 
@@ -40,7 +41,7 @@ function precompiling()
     return (@ccall jl_generating_output()::Cint) == 1
 end
 
-struct ReactantABI <: Enzyme.EnzymeCore.ABI end
+struct ReactantABI <: EnzymeCore.ABI end
 
 include("PrimitiveTypes.jl")
 

test/autodiff.jl

@@ -286,15 +286,15 @@ function simple_grad_without_ignore(x::AbstractArray{T}) where {T}
 end
 
 function simple_grad_with_ignore(x::AbstractArray{T}) where {T}
-    return Reactant.ignore_derivatives(sum(x; dims=1), x .- 1, (x, x .+ 2)), sum(abs2, x)
+    return Enzyme.ignore_derivatives(sum(x; dims=1), x .- 1, (x, x .+ 2)), sum(abs2, x)
 end
 
 function zero_grad(x)
-    return Reactant.ignore_derivatives(sum(x))
+    return Enzyme.ignore_derivatives(sum(x))
 end
 
 function zero_grad2(x)
-    return Reactant.ignore_derivatives(sum(x), x)
+    return Enzyme.ignore_derivatives(sum(x), x)
 end
 
 @testset "ignore_derivatives" begin