Skip to content

Commit 70fc3cd

Browse files
BSnellingmbauman
authored andcommitted
define _maxndims methods for small tuples to help inference
1 parent 1b6ffda commit 70fc3cd

File tree

1 file changed

+8
-4
lines changed

1 file changed

+8
-4
lines changed

base/broadcast.jl

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -263,11 +263,15 @@ end
263263

264264
Base.IteratorSize(::Type{T}) where {T<:Broadcasted} = Base.HasShape{ndims(T)}()
265265
Base.ndims(BC::Type{<:Broadcasted{<:Any,Nothing}}) = _maxndims(fieldtype(BC, 2))
266-
function Base.ndims(BC::Type{<:Broadcasted{<:AbstractArrayStyle{N},Nothing}}) where {N}
267-
N isa Integer && return N
268-
_maxndims(fieldtype(BC, 2))
266+
Base.ndims(::Type{<:Broadcasted{<:AbstractArrayStyle{N},Nothing}}) where {N<:Integer} = N
267+
268+
_maxndims(T::Type{<:Tuple}) = reduce(max, (ntuple(n -> _ndims(fieldtype(T, n)), Base._counttuple(T))))
269+
_maxndims(::Type{<:Tuple{T}}) where {T} = ndims(T)
270+
_maxndims(::Type{<:Tuple{T}}) where {T<:Tuple} = _ndims(T)
271+
function _maxndims(::Type{<:Tuple{T, S}}) where {T, S}
272+
return T<:Tuple || S<:Tuple ? max(_ndims(T), _ndims(S)) : max(ndims(T), ndims(S))
269273
end
270-
Base.@pure _maxndims(T) = mapfoldl(_ndims, max, fieldtypes(T))
274+
271275
_ndims(x) = ndims(x)
272276
_ndims(::Type{<:Tuple}) = 1
273277

0 commit comments

Comments
 (0)