fix: handle traced array returns inside objects

Author: avik-pal

src/Compiler.jl

@@ -631,17 +631,45 @@ function codegen_unflatten!(
                 path = path[2:end]
                 result_stores[path] = concrete_res_name
                 continue
-            else
-                @assert path[1] == :resargs
+            elseif path[1] == :resargs
                 unflatcode = :(args[$(path[2])])
                 path = path[3:end]
+            else
+                @show "path[1] == $(path[1]) has been ignored..." # XXX: Validate this is correct
+                continue
             end
 
             # unroll path tree
-            for p in path
+            for p in path[1:(end - 1)]
                 unflatcode = :(traced_getfield($unflatcode, $(Meta.quot(p))))
             end
-            unflatcode = :($unflatcode.data = $concrete_res_name)
+            if length(path) > 0
+                final_val = gensym("final_val")
+                unflatcode = quote
+                    $final_val = traced_getfield($unflatcode, $(Meta.quot(path[end])))
+                    if $final_val isa TracedRArray
+                        setfield!(
+                            $unflatcode,
+                            $(Meta.quot(path[end])),
+                            ConcreteRArray{eltype($final_val),ndims($final_val)}(
+                                $concrete_res_name, size($final_val)
+                            ),
+                        )
+                    elseif $final_val isa TracedRNumber
+                        setfield!(
+                            $unflatcode,
+                            $(Meta.quot(path[end])),
+                            ConcreteRNumber{eltype($final_val)}($concrete_res_name),
+                        )
+                    else
+                        setfield!($final_val, :data, $concrete_res_name)
+                    end
+                end
+            else
+                unflatcode = quote
+                    $unflatcode.data = $concrete_res_name
+                end
+            end
 
             push!(unflatten_code, unflatcode)
         end

src/Interpreter.jl

@@ -235,7 +235,7 @@ function overload_autodiff(
     primf = f.val
     primargs = ((v.val for v in args)...,)
 
-    fnwrap, func2, traced_result, result, seen_args, ret, linear_args, in_tys, linear_results = Reactant.TracedUtils.make_mlir_fn(
+    fnwrap, func2, traced_result, result, seen_args, ret, linear_args, in_tys, linear_results = TracedUtils.make_mlir_fn(
         primf, primargs, (), string(f) * "_autodiff", false
     )
 
@@ -302,7 +302,7 @@ function overload_autodiff(
                 cst = MLIR.IR.result(MLIR.Dialects.stablehlo.constant(; value=attr), 1)
                 push!(ad_inputs, cst)
             end
-        else
+        elseif TracedUtils.has_argidx(a)
             idx, path = TracedUtils.get_argidx(a)
             if idx == 1 && fnwrap
                 act = act_from_type(f, reverse, true)
@@ -322,6 +322,12 @@ function overload_autodiff(
                 end
                 TracedUtils.push_val!(ad_inputs, args[idx].dval, path[3:end])
             end
+        else
+            act = act_from_type(Enzyme.Const, reverse, true)
+            push!(ret_activity, act)
+            if act != enzyme_out && act != enzyme_outnoneed
+                continue
+            end
         end
     end
 
@@ -385,7 +391,7 @@ function overload_autodiff(
                 end
                 residx += 1
             end
-        else
+        elseif TracedUtils.has_argidx(a)
             idx, path = TracedUtils.get_argidx(a)
             if idx == 1 && fnwrap
                 TracedUtils.set!(
@@ -405,6 +411,9 @@ function overload_autodiff(
                 )
                 residx += 1
             end
+        else
+            TracedUtils.set!(a, (), TracedUtils.transpose_val(MLIR.IR.result(res, residx)))
+            residx += 1
         end
     end
 

src/TracedUtils.jl

@@ -341,6 +341,18 @@ function get_argidx(x)
     throw(AssertionError("No path found for $x"))
 end
 
+function has_argidx(x)
+    for path in x.paths
+        if length(path) == 0
+            continue
+        end
+        if path[1] == :args
+            return true
+        end
+    end
+    return false
+end
+
 function set!(x, path, tostore; emptypath=false)
     for p in path
         x = Reactant.Compiler.traced_getfield(x, p)