Refactor linear_args, linear_results

EnzymeAD · Sep 4, 2024 · 155ddb6 · 155ddb6
1 parent d44c7da
commit 155ddb6
Showing 1 changed file with 2 additions and 16 deletions.
diff --git a/src/Compiler.jl b/src/Compiler.jl
@@ -343,13 +343,7 @@ function make_mlir_fn(f, args, kwargs, name="main", concretein=true; toscalar=fa
         )
     end
 
-    linear_args = TracedRArray[]
-    for (k, v) in seen_args
-        if !(v isa TracedRArray)
-            continue
-        end
-        push!(linear_args, v)
-    end
+    linear_args = TracedRArray[v for v in values(seen_args) if v isa TracedRArray]
 
     in_tys = if toscalar
         [MLIR.IR.TensorType((), MLIR.IR.Type(eltype(arg))) for arg in linear_args]
@@ -437,15 +431,7 @@ function make_mlir_fn(f, args, kwargs, name="main", concretein=true; toscalar=fa
         )
     end
 
-    linear_results = TracedRArray[]
-
-    for (k, v) in seen_results
-        if !(v isa TracedRArray)
-            continue
-        end
-
-        push!(linear_results, v)
-    end
+    linear_results = TracedRArray[v for v in values(seen_results) if v isa TracedRArray]
 
     out_tys = [transpose_ty(mlir_type(arg)) for arg in linear_results]