feat: support Base.stack

Author: avik-pal

src/Overlay.jl

@@ -127,7 +127,7 @@ for (cT, aT, bT) in (
         @reactant_overlay @noinline function LinearAlgebra.mul!(
             C::$cT, A::$aT, B::$bT, α::Number, β::Number
         )
-            if any(Base.Fix2(isa, TracedRArray) ∘ ancestor, (C, A, B))
+            if use_overlayed_version((C, A, B))
                 TracedLinearAlgebra.overloaded_mul!(C, A, B, α, β)
             else
                 LinearAlgebra._mul!(C, A, B, α, β)
@@ -142,3 +142,12 @@ for (cT, aT, bT) in (
         end
     end
 end
+
+# Base overloads
+@reactant_overlay @noinline function Base._stack(dims::Union{Integer,Colon}, iter)
+    if use_overlayed_version(iter)
+        return TracedRArrayOverrides.overloaded_stack(dims, iter)
+    else
+        return Base._stack(dims, Base.IteratorSize(iter), iter)
+    end
+end

src/Reactant.jl

@@ -185,6 +185,20 @@ mutable struct TracedRNG <: Random.AbstractRNG
     const algorithm::String
 end
 
+use_overlayed_version(iter) = any(use_overlayed_version, iter)
+
+use_overlayed_version(::TracedRArray) = true
+use_overlayed_version(::TracedRNumber) = true
+use_overlayed_version(::Number) = false
+use_overlayed_version(::MissingTracedValue) = true
+use_overlayed_version(::TracedRNG) = true
+
+function use_overlayed_version(x::AbstractArray)
+    a = ancestor(x)
+    a === x && return false
+    return use_overlayed_version(a)
+end
+
 # StdLib Overloads
 include("stdlibs/LinearAlgebra.jl")
 include("stdlibs/Random.jl")

src/TracedRArray.jl

@@ -605,4 +605,17 @@ function Base._RepeatInnerOuter.repeat_inner(
     return materialize_traced_array(reshape(x_broadcasted, final_size...))
 end
 
+# stack
+function overloaded_stack(dims::Union{Integer,Colon}, xs)
+    @assert allequal(ndims, xs) "All arrays must have the same number of dimensions..."
+    dims = dims isa Colon ? ndims(first(xs)) + 1 : dims
+    res = map(xs) do x
+        new_shape = ntuple(
+            i -> i == dims ? 1 : (i < dims ? size(x, i) : size(x, i - 1)), ndims(x) + 1
+        )
+        return materialize_traced_array(reshape(x, new_shape))
+    end
+    return cat(res...; dims)
+end
+
 end

src/stdlibs/LinearAlgebra.jl

@@ -269,7 +269,7 @@ function LinearAlgebra.diag(x::AnyTracedRArray{T,2}, k::Integer=0) where {T}
     #   <unknown>:0: note: see current operation: %0 = "tensor.empty"() : () -> tensor<0xf64>
     length(indices) ≤ 0 && return TracedUtils.promote_to(TracedRArray{T,1}, T[])
 
-    return Ops.gather_getindex(x, TracedUtils.promote_to(TracedRArray{Int,2}, indices))
+    return Ops.gather_getindex(y, TracedUtils.promote_to(TracedRArray{Int,2}, indices))
 end
 
 function LinearAlgebra._diagm(