Skip to content

Preserve shape when collecting broadcasted objects #44061

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

Merged
merged 6 commits into from
May 27, 2022
Merged

Preserve shape when collecting broadcasted objects #44061

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
a5575a0
define IteratorSize for array style broadcasted
Feb 7, 2022
e0511da
test collected broadcasted objects retain their shape
Feb 7, 2022
f0049ba
IteratorSize for ArrayStyle broadcasts that don't propagate dims
Feb 14, 2022
81efab9
Generalise IteratorSize definition for broadcasted
Feb 17, 2022
1b6ffda
support itertor size for nested broadcasts using @pure
Feb 21, 2022
70fc3cd
define _maxndims methods for small tuples to help inference
Mar 4, 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
17 changes: 15 additions & 2 deletions base/broadcast.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -244,7 +244,7 @@ Base.IndexStyle(::Type{<:Broadcasted{<:Any}}) = IndexCartesian()

Base.LinearIndices(bc::Broadcasted{<:Any,<:Tuple{Any}}) = LinearIndices(axes(bc))::LinearIndices{1}

Base.ndims(::Broadcasted{<:Any,<:NTuple{N,Any}}) where {N} = N
Base.ndims(bc::Broadcasted) = ndims(typeof(bc))
Base.ndims(::Type{<:Broadcasted{<:Any,<:NTuple{N,Any}}}) where {N} = N

Base.size(bc::Broadcasted) = map(length, axes(bc))
Expand All @@ -261,7 +261,20 @@ Base.@propagate_inbounds function Base.iterate(bc::Broadcasted, s)
return (bc[i], (s[1], newstate))
end

Base.IteratorSize(::Type{<:Broadcasted{<:Any,<:NTuple{N,Base.OneTo}}}) where {N} = Base.HasShape{N}()
Base.IteratorSize(::Type{T}) where {T<:Broadcasted} = Base.HasShape{ndims(T)}()
Base.ndims(BC::Type{<:Broadcasted{<:Any,Nothing}}) = _maxndims(fieldtype(BC, 2))
Base.ndims(::Type{<:Broadcasted{<:AbstractArrayStyle{N},Nothing}}) where {N<:Integer} = N
Copy link
Member

@N5N3 N5N3 May 27, 2022 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks like this line will never be hitted.
So even AbstractArrayStyle with dimension tracking now use the general fallback above.
Thus we have

julia> Base.broadcasted(randn,) |> collect
ERROR: MethodError: reducing over an empty collection is not allowed; consider supplying `init` to the reducer


_maxndims(T::Type{<:Tuple}) = reduce(max, (ntuple(n -> _ndims(fieldtype(T, n)), Base._counttuple(T))))
_maxndims(::Type{<:Tuple{T}}) where {T} = ndims(T)
_maxndims(::Type{<:Tuple{T}}) where {T<:Tuple} = _ndims(T)
function _maxndims(::Type{<:Tuple{T, S}}) where {T, S}
return T<:Tuple || S<:Tuple ? max(_ndims(T), _ndims(S)) : max(ndims(T), ndims(S))
end

_ndims(x) = ndims(x)
_ndims(::Type{<:Tuple}) = 1

Base.IteratorEltype(::Type{<:Broadcasted}) = Base.EltypeUnknown()

## Instantiation fills in the "missing" fields in Broadcasted.
Expand Down
33 changes: 33 additions & 0 deletions test/broadcast.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -855,6 +855,39 @@ let
@test ndims(copy(bc)) == ndims([v for v in bc]) == ndims(collect(bc)) == ndims(bc)
end

# issue 43847: collect preserves shape of broadcasted
let
bc = Broadcast.broadcasted(*, [1 2; 3 4], 2)
@test collect(Iterators.product(bc, bc)) == collect(Iterators.product(copy(bc), copy(bc)))

a1 = AD1(rand(2,3))
bc1 = Broadcast.broadcasted(*, a1, 2)
@test collect(Iterators.product(bc1, bc1)) == collect(Iterators.product(copy(bc1), copy(bc1)))

# using ndims of second arg
bc2 = Broadcast.broadcasted(*, 2, a1)
@test collect(Iterators.product(bc2, bc2)) == collect(Iterators.product(copy(bc2), copy(bc2)))

# >2 args
bc3 = Broadcast.broadcasted(*, a1, 3, a1)
@test collect(Iterators.product(bc3, bc3)) == collect(Iterators.product(copy(bc3), copy(bc3)))

# including a tuple and custom array type
bc4 = Broadcast.broadcasted(*, (1,2,3), AD1(rand(3)))
@test collect(Iterators.product(bc4, bc4)) == collect(Iterators.product(copy(bc4), copy(bc4)))

# testing ArrayConflict
@test Broadcast.broadcasted(+, AD1(rand(3)), AD2(rand(3))) isa Broadcast.Broadcasted{Broadcast.ArrayConflict}
@test Broadcast.broadcasted(+, AD1(rand(3)), AD2(rand(3))) isa Broadcast.Broadcasted{<:Broadcast.AbstractArrayStyle{Any}}

@test @inferred(Base.IteratorSize(Broadcast.broadcasted((1,2,3),a1,zeros(3,3,3)))) === Base.HasShape{3}()

# inference on nested
bc = Base.broadcasted(+, AD1(randn(3)), AD1(randn(3)))
bc_nest = Base.broadcasted(+, bc , bc)
@test @inferred(Base.IteratorSize(bc_nest)) === Base.HasShape{1}()
end

# issue #31295
let a = rand(5), b = rand(5), c = copy(a)
view(identity(a), 1:3) .+= view(b, 1:3)
Expand Down