Permit autodiff thunk to use provided RT (#2200)

Co-authored-by: William Moses <[email protected]>
EnzymeAD · Dec 15, 2024 · 1dc7f10 · 1dc7f10
1 parent ee21090
commit 1dc7f10
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/src/Enzyme.jl b/src/Enzyme.jl
@@ -1381,10 +1381,13 @@ result, ∂v, ∂A
 
     TT = Tuple{args...}
 
-    primal_tt = Tuple{map(eltype, args)...}
-    rt0 = Compiler.primal_return_type(Reverse, eltype(FA), primal_tt)
-
-    rt = Compiler.remove_innerty(A2){rt0}
+    rt = if A2 isa UnionAll
+        primal_tt = Tuple{map(eltype, args)...}
+	rt0 = Compiler.primal_return_type(Reverse, eltype(FA), primal_tt)
+	A2{rt0}
+    else
+	A2
+    end
 
     primal_ptr = Compiler.deferred_codegen(
         FA,