From f2dec4c047615d87dd151aa06659e7796fc3297c Mon Sep 17 00:00:00 2001
From: Billy Moses <wmoses@google.com>
Date: Tue, 3 Dec 2024 08:40:49 -0600
Subject: [PATCH] fix

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

diff --git a/src/compiler.jl b/src/compiler.jl
index 10dc9842d1..360807f4eb 100644
--- a/src/compiler.jl
+++ b/src/compiler.jl
@@ -5261,7 +5261,7 @@ end
 const DumpPostOpt = Ref(false)
 
 # actual compilation
-function _thunk(job, postopt::Bool = true)
+function _thunk(job, postopt::Bool = true)::Tuple{LLVM.Module, String, Union{String, Nothing}, Type, String}
     mod, meta = codegen(:llvm, job; optimize = false)
     adjointf, augmented_primalf = meta.adjointf, meta.augmented_primalf
 
@@ -5315,15 +5315,16 @@ const cache_lock = ReentrantLock()
         if obj === nothing
             asm = _thunk(job)
             obj = _link(job, asm...)
+            @show obj
             if obj.adjoint isa Ptr{Nothing}
                 autodiff_cache[obj.adjoint] = (asm[2], asm[5])
             end
-            if obj.primal isa Ptr{Nothing}
+            if obj.primal isa Ptr{Nothing} && asm[3] isa String
                 autodiff_cache[obj.primal] = (asm[3], asm[5])
             end
             cache[key] = obj
         end
-        nothing
+        obj
     finally
         unlock(cache_lock)
     end