feat: error hints for Enzyme (#788)

* feat: error hints for Enzyme

* Clearer message

* Add tests

* Simpler show

* Unreleased

* Move

Author: gdalle

DifferentiationInterface/CHANGELOG.md

@@ -7,13 +7,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
-### Changed
+### Added
 
-- Allocate Enzyme shadow memory during preparation ([#782])
+- Error hints for Enzyme ([#788])
 
 ## [0.6.53] - 2025-05-07
 
+### Changed
+
+- Allocate Enzyme shadow memory during preparation ([#782])
+
 [unreleased]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.6.53...main
 [0.6.53]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.6.52...DifferentiationInterface-v0.6.53
 
+[#788]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/788
 [#782]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/782

DifferentiationInterface/Project.toml

@@ -1,7 +1,7 @@
 name = "DifferentiationInterface"
 uuid = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 authors = ["Guillaume Dalle", "Adrian Hill"]
-version = "0.6.53"
+version = "0.6.54"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -56,7 +56,7 @@ ADTypes = "1.13.0"
 ChainRulesCore = "1.23.0"
 DiffResults = "1.1.0"
 Diffractor = "=0.2.6"
-Enzyme = "0.13.17"
+Enzyme = "0.13.39"
 EnzymeCore = "0.8.8"
 ExplicitImports = "1.10.1"
 FastDifferentiation = "0.4.3"

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/DifferentiationInterfaceEnzymeExt.jl

@@ -30,6 +30,7 @@ using EnzymeCore:
     Split,
     WithPrimal
 using Enzyme:
+    Enzyme,
     autodiff,
     autodiff_thunk,
     create_shadows,
@@ -53,4 +54,6 @@ include("forward_twoarg.jl")
 include("reverse_onearg.jl")
 include("reverse_twoarg.jl")
 
+include("init.jl")
+
 end # module

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/init.jl

@@ -0,0 +1,44 @@
+const HINT_END = "\n\nThis hint appears because DifferentiationInterface and Enzyme are both loaded. It does not necessarily imply that Enzyme is being called through DifferentiationInterface.\n\n"
+
+function HINT_START(option)
+    return "\nIf you are using Enzyme by selecting the `AutoEnzyme` object from ADTypes, you may want to try setting the `$option` option as follows:"
+end
+
+function __init__()
+    # robust against internal changes
+    condition = (
+        isdefined(Enzyme, :Compiler) &&
+        Enzyme.Compiler isa Module &&
+        isdefined(Enzyme.Compiler, :EnzymeError) &&
+        Enzyme.Compiler.EnzymeError isa DataType
+    )
+    condition || return nothing
+    # see https://github.com/JuliaLang/julia/issues/58367 for why this isn't easier
+    for n in names(Enzyme.Compiler; all=true)
+        T = getfield(Enzyme.Compiler, n)
+        if T isa DataType && T <: Enzyme.Compiler.EnzymeError
+            # robust against internal changes
+            Base.Experimental.register_error_hint(T) do io, exc
+                if occursin("EnzymeMutabilityException", string(nameof(T)))
+                    printstyled(io, HINT_START("function_annotation"); bold=true)
+                    printstyled(
+                        io,
+                        "\n\n\tAutoEnzyme(; function_annotation=Enzyme.Duplicated)";
+                        color=:cyan,
+                        bold=true,
+                    )
+                    printstyled(io, HINT_END; italic=true)
+                elseif occursin("EnzymeRuntimeActivityError", string(nameof(T)))
+                    printstyled(io, HINT_START("mode"); bold=true)
+                    printstyled(
+                        io,
+                        "\n\n\tAutoEnzyme(; mode=Enzyme.set_runtime_activity(Enzyme.Forward))\n\tAutoEnzyme(; mode=Enzyme.set_runtime_activity(Enzyme.Reverse))";
+                        color=:cyan,
+                        bold=true,
+                    )
+                    printstyled(io, HINT_END; italic=true)
+                end
+            end
+        end
+    end
+end

DifferentiationInterface/test/Back/Enzyme/test.jl

@@ -148,3 +148,56 @@ end
         f_nocontext, AutoEnzyme(; mode=Enzyme.Reverse), rand(10), ConstantOrCache(nothing)
     )
 end
+
+@testset "Hints" begin
+    @testset "MutabilityError" begin
+        f = let
+            cache = [0.0]
+            x -> sum(copyto!(cache, x))
+        end
+
+        e = nothing
+        try
+            gradient(f, AutoEnzyme(), [1.0])
+        catch e
+        end
+        msg = sprint(showerror, e)
+        @test occursin("AutoEnzyme", msg)
+        @test occursin("function_annotation", msg)
+        @test occursin("ADTypes", msg)
+        @test occursin("DifferentiationInterface", msg)
+    end
+
+    @testset "RuntimeActivityError" begin
+        function g(active_var, constant_var, cond)
+            if cond
+                return active_var
+            else
+                return constant_var
+            end
+        end
+
+        function h(active_var, constant_var, cond)
+            return [g(active_var, constant_var, cond), g(active_var, constant_var, cond)]
+        end
+
+        e = nothing
+        try
+            pushforward(
+                h,
+                AutoEnzyme(; mode=Enzyme.Forward),
+                [1.0],
+                ([1.0],),
+                Constant([1.0]),
+                Constant(true),
+            )
+        catch e
+        end
+        msg = sprint(showerror, e)
+        @test occursin("AutoEnzyme", msg)
+        @test occursin("mode", msg)
+        @test occursin("set_runtime_activity", msg)
+        @test occursin("ADTypes", msg)
+        @test occursin("DifferentiationInterface", msg)
+    end
+end

DifferentiationInterface/test/Core/Internals/hints.jl

@@ -0,0 +1,14 @@
+using ADTypes
+using DifferentiationInterface
+import DifferentiationInterface as DI
+using Test
+
+@testset "Missing backend" begin
+    e = nothing
+    try
+        gradient(sum, AutoZygote(), [1.0])
+    catch e
+    end
+    msg = sprint(showerror, e)
+    @test occursin("import Zygote", msg)
+end