Fix size implementation for MockTensor tests

EnzymeAD · Aug 13, 2024 · 175f538 · 175f538
1 parent 69ebd41
commit 175f538
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/test/struct.jl b/test/struct.jl
@@ -2,17 +2,18 @@ using Reactant
 using Test
 
 # from bsc-quantic/Tenet.jl
-struct MockTensor{T,N,A<:AbstractArray{T,N}}
+struct MockTensor{T,N,A<:AbstractArray{T,N}} <: AbstractArray{T,N}
     data::A
     inds::Vector{Symbol}
 end
 
 MockTensor(data::A, inds) where {T,N,A<:AbstractArray{T,N}} = MockTensor{T,N,A}(data, inds)
 Base.parent(t::MockTensor) = t.data
+Base.size(t::MockTensor) = size(parent(t))
 
 Base.cos(x::MockTensor) = MockTensor(cos(parent(x)), x.inds)
 
-mutable struct MutableMockTensor{T,N,A<:AbstractArray{T,N}}
+mutable struct MutableMockTensor{T,N,A<:AbstractArray{T,N}} <: AbstractArray{T,N}
     data::A
     inds::Vector{Symbol}
 end
@@ -21,6 +22,7 @@ function MutableMockTensor(data::A, inds) where {T,N,A<:AbstractArray{T,N}}
     return MutableMockTensor{T,N,A}(data, inds)
 end
 Base.parent(t::MutableMockTensor) = t.data
+Base.size(t::MutableMockTensor) = size(parent(t))
 
 Base.cos(x::MutableMockTensor) = MutableMockTensor(cos(parent(x)), x.inds)