refactor: use a union type for traced types

Author: avik-pal

src/Compiler.jl

@@ -290,7 +290,7 @@ function compile_mlir!(mod, f, args; optimize=true)
     preserved_args = Tuple{TracedRArray,Int}[]
     results = [MLIR.IR.operand(ret, i) for i in 1:MLIR.IR.noperands(ret)]
     nresults = MLIR.IR.Value[]
-    linear_results2 = Union{TracedRArray,TracedRNumber}[]
+    linear_results2 = TracedTypes[]
     for (i, op) in enumerate(results)
         if !MLIR.IR.is_block_arg(op)
             push!(nresults, op)

src/Reactant.jl

@@ -83,9 +83,13 @@ include("mlir/MLIR.jl")
 include("XLA.jl")
 include("Interpreter.jl")
 include("utils.jl")
+
 include("ConcreteRArray.jl")
 include("TracedRNumber.jl")
 include("TracedRArray.jl")
+
+const TracedTypes = Union{TracedRArray,TracedRNumber}
+
 include("Tracing.jl")
 include("Compiler.jl")
 

src/utils.jl

@@ -44,7 +44,7 @@ function make_mlir_fn(f, args, kwargs, name="main", concretein=true; toscalar=fa
         )
     end
 
-    linear_args = Union{TracedRArray,TracedRNumber}[]
+    linear_args = TracedTypes[]
     for (k, v) in seen_args
         if !(v isa TracedRArray) && !(v isa TracedRNumber)
             continue
@@ -127,7 +127,7 @@ function make_mlir_fn(f, args, kwargs, name="main", concretein=true; toscalar=fa
         )
     end
 
-    linear_results = Union{TracedRArray,TracedRNumber}[]
+    linear_results = TracedTypes[]
 
     for (k, v) in seen_results
         if !(v isa TracedRArray) && !(v isa TracedRNumber)