fix: ensure printing of wrapped ConcreteRArrays goes through our show (EnzymeAD#367)

* fix: ensure printing of wrapped ConcreteRArrays goes through our show

* fix: allow wrapped arrays in mapreduce

Author: avik-pal

src/ConcreteRArray.jl

@@ -195,16 +195,18 @@ function Base.show(io::IO, X::ConcreteRScalar{T}) where {T}
     return nothing
 end
 
-function Base.print_array(io::IO, X::ConcreteRArray)
-    if X.data == XLA.AsyncEmptyBuffer
+function Base.print_array(io::IO, X::AnyConcreteRArray)
+    data = ancestor(X).data
+    if data == XLA.AsyncEmptyBuffer
         println(io, "<Empty buffer>")
         return nothing
     end
     return Base.print_array(io, convert(Array, X))
 end
 
-function Base.show(io::IO, X::ConcreteRArray)
-    if X.data == XLA.AsyncEmptyBuffer
+function Base.show(io::IO, X::AnyConcreteRArray)
+    data = ancestor(X).data
+    if data == XLA.AsyncEmptyBuffer
         println(io, "<Empty buffer>")
         return nothing
     end

src/Reactant.jl

@@ -85,6 +85,12 @@ function Enzyme.make_zero(
     return res
 end
 
+function ancestor(x::AbstractArray)
+    p_x = parent(x)
+    p_x === x && return x
+    return ancestor(p_x)
+end
+
 include("mlir/MLIR.jl")
 include("XLA.jl")
 include("Interpreter.jl")

src/TracedRArray.jl

@@ -122,9 +122,6 @@ function set_mlir_data!(x::AnyTracedRArray, data)
     return x
 end
 
-ancestor(x::TracedRArray) = x
-ancestor(x::WrappedTracedRArray) = ancestor(parent(x))
-
 get_ancestor_indices(::TracedRArray, indices...) = indices
 function get_ancestor_indices(x::WrappedTracedRArray, indices...)
     return get_ancestor_indices(parent(x), Base.reindex(parentindices(x), indices)...)
@@ -388,10 +385,12 @@ end
 function Base.mapreduce(
     @nospecialize(f),
     @nospecialize(op),
-    @nospecialize(A::TracedRArray{T,N});
+    @nospecialize(A::AnyTracedRArray{T,N});
     dims=:,
     init=nothing,
 ) where {T,N}
+    A = materialize_traced_array(A)
+
     if dims isa Int
         dims = [dims]
     end