From fbc05f7934d5928a1545399d1e1b136caf88fdf1 Mon Sep 17 00:00:00 2001
From: William Moses <wsmoses@cyclops.juliacomputing.io>
Date: Sun, 15 Dec 2024 00:28:40 -0500
Subject: [PATCH] Permit autodiff thunk to use provided RT

---
 src/Enzyme.jl | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/src/Enzyme.jl b/src/Enzyme.jl
index 92ec9a623f..4c4f67de64 100644
--- 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,