Skip to content

Commit

Permalink
Process Fill broadcasting before others
Browse files Browse the repository at this point in the history
  • Loading branch information
jishnub committed Aug 26, 2024
1 parent 2f61d4b commit c141ff8
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 6 deletions.
20 changes: 14 additions & 6 deletions src/fillbroadcast.jl
Original file line number Diff line number Diff line change
Expand Up @@ -118,23 +118,31 @@ _isfill(f::Number) = true
_isfill(f::Ref) = true
_isfill(::Any) = false

function Base.copy(bc::Broadcast.Broadcasted{<:AbstractFillStyle{N}}) where {N}
_broadcast_maybecopy(bc::Broadcast.Broadcasted{<:AbstractFillStyle}) = copy(bc)
_broadcast_maybecopy(x) = x

function _fallback_copy(bc)
# treat the fill components
bc2 = Base.broadcasted(bc.f, map(_broadcast_maybecopy, bc.args)...)
# fallback style
S = Broadcast.Broadcasted{Broadcast.DefaultArrayStyle{ndims(bc)}}
copy(convert(S, bc2))
end

function Base.copy(bc::Broadcast.Broadcasted{<:AbstractFillStyle})
if _iszeros(bc)
return Zeros(typeof(_getindex_value(bc)), axes(bc))
elseif _isones(bc)
return Ones(typeof(_getindex_value(bc)), axes(bc))
elseif _isfill(bc)
return Fill(_getindex_value(bc), axes(bc))
else
# fallback style
S = Broadcast.Broadcasted{Broadcast.DefaultArrayStyle{N}}
copy(convert(S, bc))
_fallback_copy(bc)
end
end
# make the zero-dimensional case consistent with Base
function Base.copy(bc::Broadcast.Broadcasted{<:AbstractFillStyle{0}})
S = Broadcast.Broadcasted{Broadcast.DefaultArrayStyle{0}}
copy(convert(S, bc))
_fallback_copy(bc)
end

# some cases that preserve 0d
Expand Down
5 changes: 5 additions & 0 deletions test/runtests.jl
Original file line number Diff line number Diff line change
Expand Up @@ -1210,6 +1210,11 @@ end
F = Fill(1, 2)
@test g.(F, "a") === f.(F)
end

@testset "early binding" begin
A = ones(2) .+ (x -> rand()).(Fill(2,2))
@test all(==(A[1]), A)
end
end

@testset "map" begin
Expand Down

0 comments on commit c141ff8

Please sign in to comment.