Permit autodiff thunk to use provided RT

EnzymeAD · Dec 15, 2024 · fbc05f7 · fbc05f7
1 parent 95c3fcc
commit fbc05f7
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(mode, 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(mode, eltype(FA), primal_tt)
+	A2{rt0}
+    else
+	A2
+    end
 
     primal_ptr = Compiler.deferred_codegen(
         FA,