Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Make StructArrayStyle track inputs dimension #211

Closed
wants to merge 3 commits into from
Closed

Make StructArrayStyle track inputs dimension #211

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
cb3b187
Make `StructArrayStyle` track inputs dimension
N5N3 Feb 9, 2022
7a8fbf4
Add test for unstable broadcast
N5N3 Feb 20, 2022
a5e68f4
fix test
N5N3 Feb 20, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 11 additions & 3 deletions src/structarray.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -445,7 +445,15 @@ end
# broadcast
import Base.Broadcast: BroadcastStyle, ArrayStyle, AbstractArrayStyle, Broadcasted, DefaultArrayStyle

struct StructArrayStyle{Style} <: AbstractArrayStyle{Any} end
struct StructArrayStyle{S,N} <: AbstractArrayStyle{N} end

# Here we define the dimension tracking behaviour of StructArrayStyle
function StructArrayStyle{S,M}(::Val{N}) where {S,M,N}
if S <: AbstractArrayStyle{M}
return StructArrayStyle{typeof(S(Val(N))),N}()
end
return StructArrayStyle{S,N}()
end

@inline combine_style_types(::Type{A}, args...) where A<:AbstractArray =
combine_style_types(BroadcastStyle(A), args...)
Expand All @@ -455,9 +463,9 @@ combine_style_types(s::BroadcastStyle) = s

Base.@pure cst(::Type{SA}) where SA = combine_style_types(array_types(SA).parameters...)

BroadcastStyle(::Type{SA}) where SA<:StructArray = StructArrayStyle{typeof(cst(SA))}()
BroadcastStyle(::Type{SA}) where SA<:StructArray = StructArrayStyle{typeof(cst(SA)),ndims(SA)}()

Base.similar(bc::Broadcasted{StructArrayStyle{S}}, ::Type{ElType}) where {S<:DefaultArrayStyle,N,ElType} =
Base.similar(bc::Broadcasted{<:StructArrayStyle{S}}, ::Type{ElType}) where {S<:DefaultArrayStyle,ElType} =
isstructtype(ElType) ? similar(StructArray{ElType}, axes(bc)) : similar(Array{ElType}, axes(bc))

# for aliasing analysis during broadcast
Expand Down
19 changes: 18 additions & 1 deletion test/runtests.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -926,8 +926,25 @@ Base.similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{MyArray}}, ::Type{El
# used inside of broadcast but we also test it here explicitly
@test isa(@inferred(Base.dataids(s)), NTuple{N, UInt} where {N})

s = StructArray{ComplexF64}((MyArray(rand(2,2)), MyArray(rand(2,2))))
s = StructArray{ComplexF64}((MyArray(rand(2)), MyArray(rand(2))))
@test_throws MethodError s .+ s

# test for dimensionality track
@test Base.broadcasted(+, s, s) isa Broadcast.Broadcasted{<:Broadcast.AbstractArrayStyle{1}}
@test Base.broadcasted(+, s, 1:2) isa Broadcast.Broadcasted{<:Broadcast.AbstractArrayStyle{1}}
@test Base.broadcasted(+, s, reshape(1:2,1,2)) isa Broadcast.Broadcasted{<:Broadcast.AbstractArrayStyle{2}}
@test Base.broadcasted(+, reshape(1:2,1,1,2), s) isa Broadcast.Broadcasted{<:Broadcast.AbstractArrayStyle{3}}

a = StructArray([1;2+im])
b = StructArray([1;;2+im])
@test a .+ b == a .+ collect(b) == collect(a) .+ b == collect(a) .+ collect(b)
@test a .+ Any[1] isa StructArray

# issue #185
A = StructArray(randn(ComplexF64, 3, 3))
B = randn(ComplexF64, 3, 3)
c = StructArray(randn(ComplexF64, 3))
@test (A .= B .* c) === A
end

@testset "staticarrays" begin
Expand Down