Fix error exception (#1931)

* Fix error exception * fix --------- Co-authored-by: William Moses <[email protected]>
EnzymeAD · Sep 30, 2024 · 66ef0f3 · 66ef0f3
1 parent d97bb83
commit 66ef0f3
Showing 1 changed file with 8 additions and 4 deletions.
diff --git a/src/Enzyme.jl b/src/Enzyme.jl
@@ -682,15 +682,19 @@ code, as well as high-order differentiation.
 
     if A isa UnionAll
         rt = Compiler.primal_return_type(rmode, Val(world), FTy, tt)
-        A2 = A{rt}
+        rt = Core.Compiler.return_type(f.val, tt)
+	A2 = A{rt}
+	if rt == Union{}
+	    throw(ErrorException("Return type inferred to be Union{}. Giving up."))
+        end
     else
         @assert A isa DataType
         rt = A
+	if rt == Union{}
+	    throw(ErrorException("Return type inferred to be Union{}. Giving up."))
+        end
     end
 
-    if rt == Union{}
-        error("Return type inferred to be Union{}. Giving up.")
-    end
 
     ModifiedBetweenT = falses_from_args(Nargs + 1)
     ModifiedBetween = Val(ModifiedBetweenT)