Fix anytracedarray override (#1290)

* Fix anytracedarray override

* fix

* fix2

* add test

Author: wsmoses

src/TracedUtils.jl

@@ -22,8 +22,6 @@ ReactantCore.materialize_traced_array(x::AbstractArray) = x
 
 ReactantCore.materialize_traced_array(x::TracedRArray) = x
 
-ReactantCore.materialize_traced_array(x::AnyTracedRArray) = x[axes(x)...]
-
 function ReactantCore.materialize_traced_array(x::AbstractRange)
     return Reactant.aos_to_soa(collect(x))
 end
@@ -54,7 +52,11 @@ function ReactantCore.materialize_traced_array(
 end
 
 function ReactantCore.materialize_traced_array(x::AbstractArray{TracedRNumber{T}}) where {T}
-    return Reactant.aos_to_soa(x)
+    as = Reactant.aos_to_soa(x)
+    if as === x
+        as = x[axes(x)...]
+    end
+    return ReactantCore.materialize_traced_array(as)
 end
 
 get_mlir_data(x::TracedRNumber) = x.mlir_data

src/Types.jl

@@ -74,7 +74,7 @@ end
 @leaf TracedRArray
 Adapt.parent_type(::Type{TracedRArray{T,N}}) where {T,N} = TracedRArray{T,N}
 
-const AnyTracedRArray{T,N} = AbstractArray{<:TracedRNumber{T},N}
+const AnyTracedRArray{T,N} = AbstractArray{TracedRNumber{T},N}
 const AnyTracedRVector{T} = AnyTracedRArray{T,1}
 const AnyTracedRMatrix{T} = AnyTracedRArray{T,2}
 const AnyTracedRVecOrMat{T} = Union{AnyTracedRVector{T},AnyTracedRMatrix{T}}

test/tracing.jl

@@ -26,6 +26,10 @@ struct RMSProp{Teta,Trho,Teps,C<:Bool}
     centred::C
 end
 
+@testset "Traced Type" begin
+    @test !(Vector{Union{}} <: Reactant.AnyTracedRArray)
+end
+
 @testset "Tracing" begin
     @testset "trace_type" begin
         @testset "mode = ConcreteToTraced" begin