Use broadcasting to construct the TracedRArray

EnzymeAD · May 27, 2024 · 08d36f8 · 08d36f8
1 parent f23abb5
commit 08d36f8
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/Reactant.jl b/src/Reactant.jl
@@ -8,6 +8,7 @@ abstract type RArray{ElType,Shape,N} <: AbstractArray{ElType, N} end
 
 @inline Base.eltype(::RArray{ElType,Shape}) where {ElType, Shape} = ElType
 @inline Base.size(::RArray{ElType,Shape}) where {ElType, Shape} = Shape
+@inline Base.size(::Type{<:RArray{ElType,Shape}}) where {ElType, Shape} = Shape
 @inline Base.ndims(::RArray{ElType,Shape, N}) where {ElType, Shape, N} = N
 @inline Base.ndims(::Type{<:RArray{ElType,Shape, N}}) where {ElType, Shape, N} = N
 
@@ -140,8 +141,7 @@ end
     end
 
     if RT <: TracedRArray
-        @assert typeof(prev) <: RT # prev has concrete shape while RT might not have that
-        res = typeof(prev)(prev.paths, prev.mlir_data)
+        res = broadcast_to_size(eltype(RT)(0), size(prev))
         seen[prev] = res
         return res
     end