diff --git a/src/ConcreteRArray.jl b/src/ConcreteRArray.jl index ceb08440..dac67bf6 100644 --- a/src/ConcreteRArray.jl +++ b/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, "") 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, "") return nothing end diff --git a/src/Reactant.jl b/src/Reactant.jl index 06fd59af..0fc900b2 100644 --- a/src/Reactant.jl +++ b/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") diff --git a/src/TracedRArray.jl b/src/TracedRArray.jl index 0d73b109..5bc3ee30 100644 --- a/src/TracedRArray.jl +++ b/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