Merge pull request #31369 from JuliaLang/kf/broadcastedsize

Defined size(::Broadcasted)

Author: mbauman

base/broadcast.jl

@@ -222,7 +222,8 @@ _eachindex(t::Tuple) = CartesianIndices(t)
 Base.ndims(::Broadcasted{<:Any,<:NTuple{N,Any}}) where {N} = N
 Base.ndims(::Type{<:Broadcasted{<:Any,<:NTuple{N,Any}}}) where {N} = N
 
-Base.length(bc::Broadcasted) = prod(map(length, axes(bc)))
+Base.size(bc::Broadcasted) = map(length, axes(bc))
+Base.length(bc::Broadcasted) = prod(size(bc))
 
 function Base.iterate(bc::Broadcasted)
     iter = eachindex(bc)

test/broadcast.jl

@@ -433,7 +433,7 @@ Base.getindex(A::ArrayData, i::Integer...) = A.data[i...]
 Base.setindex!(A::ArrayData, v::Any, i::Integer...) = setindex!(A.data, v, i...)
 Base.size(A::ArrayData) = size(A.data)
 Base.similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{A}}, ::Type{T}) where {A,T} =
-    A(Array{T}(undef, length.(axes(bc))))
+    A(Array{T}(undef, size(bc)))
 
 struct Array19745{T,N} <: ArrayData{T,N}
     data::Array{T,N}
@@ -497,7 +497,7 @@ struct AD2DimStyle <: Broadcast.AbstractArrayStyle{2}; end
 AD2DimStyle(::Val{2}) = AD2DimStyle()
 AD2DimStyle(::Val{N}) where {N} = Broadcast.DefaultArrayStyle{N}()
 Base.similar(bc::Broadcast.Broadcasted{AD2DimStyle}, ::Type{T}) where {T} =
-    AD2Dim(Array{T}(undef, length.(axes(bc))))
+    AD2Dim(Array{T}(undef, size(bc)))
 Base.BroadcastStyle(::Type{T}) where {T<:AD2Dim} = AD2DimStyle()
 
 @testset "broadcasting for custom AbstractArray" begin