fix: ensure printing of wrapped ConcreteRArrays goes through our show

EnzymeAD · Dec 12, 2024 · d7d5ae6 · d7d5ae6
1 parent ea97be3
commit d7d5ae6
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 7 deletions.
diff --git 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, "<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

diff --git 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
@@ -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)...)