Skip to content

Commit 6f0a9e5

Browse files
authored
broadcast: simplify ndims (#50695)
The extra dispatch was inconsistently used, which seemed unnecessary. Also make sure T is captured as a Type parameter in the closure (fast), instead of a value (potentially slow). Very small difference measured (probably just noise): broadcast (1) | 115.68 | 0.78 | 0.7 | 4245.13 | 849.62 before broadcast (1) | 109.73 | 0.80 | 0.7 | 4243.89 | 788.25 after
1 parent a712455 commit 6f0a9e5

File tree

2 files changed

+3
-10
lines changed

2 files changed

+3
-10
lines changed

base/broadcast.jl

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -277,16 +277,12 @@ Base.IteratorSize(::Type{T}) where {T<:Broadcasted} = Base.HasShape{ndims(T)}()
277277
Base.ndims(BC::Type{<:Broadcasted{<:Any,Nothing}}) = _maxndims(fieldtype(BC, :args))
278278
Base.ndims(::Type{<:Broadcasted{<:AbstractArrayStyle{N},Nothing}}) where {N<:Integer} = N
279279

280-
_maxndims(T::Type{<:Tuple}) = reduce(max, (ntuple(n -> _ndims(fieldtype(T, n)), Base._counttuple(T))))
281-
_maxndims(::Type{<:Tuple{T}}) where {T} = ndims(T)
282-
_maxndims(::Type{<:Tuple{T}}) where {T<:Tuple} = _ndims(T)
280+
_maxndims(::Type{T}) where {T<:Tuple} = reduce(max, ntuple(n -> (F = fieldtype(T, n); F <: Tuple ? 1 : ndims(F)), Base._counttuple(T)))
281+
_maxndims(::Type{<:Tuple{T}}) where {T} = T <: Tuple ? 1 : ndims(T)
283282
function _maxndims(::Type{<:Tuple{T, S}}) where {T, S}
284-
return T<:Tuple || S<:Tuple ? max(_ndims(T), _ndims(S)) : max(ndims(T), ndims(S))
283+
return max(T <: Tuple ? 1 : ndims(T), S <: Tuple ? 1 : ndims(S))
285284
end
286285

287-
_ndims(x) = ndims(x)
288-
_ndims(::Type{<:Tuple}) = 1
289-
290286
Base.IteratorEltype(::Type{<:Broadcasted}) = Base.EltypeUnknown()
291287

292288
## Instantiation fills in the "missing" fields in Broadcasted.

test/subarray.jl

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -275,9 +275,6 @@ end
275275
# with the exception of Int-slicing
276276
oindex = (:, 6, 3:7, reshape([12]), [8,4,6,12,5,7], [3:7 1:5 2:6 4:8 5:9], reshape(2:11, 2, 5))
277277

278-
_ndims(::AbstractArray{T,N}) where {T,N} = N
279-
_ndims(x) = 1
280-
281278
if testfull
282279
let B = copy(reshape(1:13^3, 13, 13, 13))
283280
@testset "full tests: ($o1,$o2,$o3)" for o3 in oindex, o2 in oindex, o1 in oindex

0 commit comments

Comments
 (0)