Add non scalar return error (#2404)

* Add non scalar return error

* fix

* Update errors.jl

Author: wsmoses

src/compiler.jl

@@ -4986,10 +4986,7 @@ function add_one_in_place(x)
     elseif x isa (Array{T,0} where T)
         x[] = recursive_add(x[], default_adjoint(eltype(Core.Typeof(x))))
     else
-        error(
-            "Enzyme Mutability Error: Cannot add one in place to immutable value " *
-            string(x),
-        )
+        throw(EnzymeNonScalarReturnException(x, ""))
     end
     return nothing
 end

src/errors.jl

@@ -826,3 +826,20 @@ end
     end
     throw(AssertionError("Unknown errtype"))
 end
+
+struct EnzymeNonScalarReturnException <: EnzymeError
+    object
+    extra::String
+end
+
+function Base.showerror(io::IO, ece::EnzymeNonScalarReturnException)
+    if isdefined(Base.Experimental, :show_error_hints)
+        Base.Experimental.show_error_hints(io, ece)
+    end
+    println(io, "Return type of differentiated function was not a scalar as required, found ", ece.object)
+    println(io, "If calling Enzyme.autodiff(Reverse, f, Active, ...), try Enzyme.autodiff_thunk(Reverse, f, Duplicated, ....)")
+    println(io, "If calling Enzyme.gradient, try Enzyme.jacobian")
+    if length(ece.extra) != 0
+        print(io, ece.extra)
+    end
+end

src/rules/jitrules.jl

@@ -606,18 +606,15 @@ function body_runtime_generic_rev(N, Width, wrapped, primttypes, shadowargs, act
                 elseif $shad isa Base.RefValue
                     $shad[] = recursive_add($shad[], $expr)
                 else
-                    error(
-                        "Enzyme Mutability Error: Cannot add one in place to immutable value " *
-                        string($shad) *
-                        " tup[i]=" *
+                    throw(EnzymeNonScalarReturnException($shad, " tup[i]=" *
                         string(tup[$i]) *
                         " i=" *
                         string($i) *
                         " w=" *
                         string($w) *
                         " tup=" *
                         string(tup),
-                    )
+                    ))
                 end
             end
             @inbounds outs[(i-1)*Width+w] = out

src/rules/typeunstablerules.jl

@@ -124,10 +124,7 @@ function body_construct_rev(
                     if $shad isa Base.RefValue
                         $shad[] = recursive_add($shad[], $expr, identity, guaranteed_nonactive)
                     else
-                        error(
-                            "Enzyme Mutability Error: Cannot add one in place to immutable value " *
-                            string($shad),
-                        )
+                        throw(EnzymeNonScalarReturnException($shad, ""))
                     end
                 end
             )

test/errors.jl

@@ -0,0 +1,7 @@
+using Enzyme, Test
+
+array_square(x) = 2 .* x
+
+@testset "Array of Pointer Copy" begin
+	@test_throws Enzyme.Compiler.EnzymeNonScalarReturnException Enzyme.gradient(Reverse, array_square, [2.0])
+end

test/runtests.jl

@@ -2968,6 +2968,7 @@ end
 end
 
 include("sugar.jl")
+include("errors.jl")
 
 @testset "Forward on Reverse" begin