Skip to content

Fix wrong results in sparse prod, mapreduce and findmin, and make map!/broadcast! safe with missing, fixed destinations and a throwing f #1022

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

Closed
ViralBShah wants to merge 3 commits into from
Closed

Fix wrong results in sparse prod, mapreduce and findmin, and make map!/broadcast! safe with missing, fixed destinations and a throwing f #1022

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
6c54a66
Make sparse `map`, `broadcast` and reductions agree with dense on non…
ViralBShah Oct 9, 2026
ab82a14
Check every predicate value in sparse `any` and `all`
ViralBShah Oct 9, 2026
0d23e20
Merge remote-tracking branch 'origin/main' into vs/broadcast-reductio…
ViralBShah Oct 13, 2026
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
138 changes: 104 additions & 34 deletions src/higherorderfns.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,8 @@ using ..SparseArrays: SparseVector, SparseMatrixCSC, FixedSparseCSC, SparseMatri
SparseMatrixCSCColumnSubset, SparseColumnView, SparseVectorView, SparseVectorPartialView,
indtype, fixed, move_fixed, nnz, nzrange, spzeros, ftranspose,
nonzeroinds, nonzeros, rowvals, getcolptr, getrowval, getnzval, widelength,
_iszero, _isnotzero, _is_fixed, _checkbuffers, _densestructure!, _copyto_fixed!,
_iszero, _isnotzero, _is_fixed, _checkbuffers, _densestructure!, _checkdensestructure,
_copyto_fixed!, _densezero,
ReadOnly, _storagebuffers, _unaliasedcopy
using Base.Broadcast: BroadcastStyle, Broadcasted, flatten
using LinearAlgebra
Expand Down Expand Up @@ -197,21 +198,62 @@ _aliasesbuffers(Cbuffers, b) = isempty(b) ? any(c -> c === b, Cbuffers) : any(c
# Avoid calculating f(zero) unless necessary as it may fail.
let r = if _haszeros(A) && all(_haszeros, Bs)
fofzeros = f(_zeros_eltypes(A, Bs...)...)
_iszero(fofzeros) ? _map_zeropres!(f, C, A, Bs...) :
_map_notzeropres!(f, fofzeros, C, A, Bs...)
_iszeroresult(fofzeros) ? _guardpattern(_map_zeropres!, f, C, A, Bs...) :
_guardfill(_map_notzeropres!, f, fofzeros, C, A, Bs...)
else
_map_zeropres!(f, C, A, Bs...)
_guardpattern(_map_zeropres!, f, C, A, Bs...)
end
return r
end
# A zero-preserving kernel rebuilds the pattern of C as it goes. When `f` or the conversion
# of its result throws part-way, a C that had the pattern of the result is left partly
# written, as a dense destination is, and any other C structurally empty rather than with a
# pattern that is half its own. A fixed C keeps its pattern throughout.
function _guardpattern(kernel!::K, f::Tf, C::SparseVecOrMat, As::Vararg{SparseVecOrMat,N}) where {K,Tf,N}
try
return kernel!(f, C, As...)
catch
_is_fixed(C) || _keepspattern(C, As...) || _emptypattern!(C)
rethrow()
end
end
# The result of one argument of the shape of C has the pattern of that argument, so C held
# it all along if it still does. The pattern of another result depends on its values.
_keepspattern(C, A) = _aresameshape(C, A) && _samepattern(C, A)
_keepspattern(C, As...) = false
_emptypattern!(C) = (trimstorage!(C, 0); _finishempty!(C))
# A kernel that fills C stores every entry of it first. What can be checked is checked before
# that, the conversion of the fill value included, and leaves C as it was. A later failure
# leaves a C that is not fixed structurally empty, as its values are no longer its own.
function _guardfill(kernel!::K, f::Tf, fillvalue, C::SparseVecOrMat, As::Vararg{SparseVecOrMat,N}) where {K,Tf,N}
_checkdensestructure(C)
convert(eltype(C), fillvalue)
try
return kernel!(f, fillvalue, C, As...)
catch
_is_fixed(C) || _emptypattern!(C)
rethrow()
end
end
# Whether `f` at the zeros is a zero. A value of a type that has none, such as a `String`,
# a `Tuple` or `nothing`, is not one, and fills the result like any other nonzero.
_iszeroresult(x::Union{Number,AbstractArray,Missing}) = _iszero(x)
function _iszeroresult(x)
try
return _iszero(x)
catch e
e isa MethodError || rethrow()
return false
end
end

function _noshapecheck_map(f::Tf, A::SparseVecOrMat, Bs::Vararg{SparseVecOrMat,N}) where {Tf,N}
# Avoid calculating f(zero) unless necessary as it may fail.
entrytypeC = Base.promote_typejoin_union(Base.promote_op(f, map(eltype, (A, Bs...))...))
indextypeC = _promote_indtype(A, Bs...)
r = if _haszeros(A) && all(_haszeros, Bs)
fofzeros = f(_zeros_eltypes(A, Bs...)...)
fpreszeros = _iszero(fofzeros)
fpreszeros = _iszeroresult(fofzeros)
maxnnzC = Int(fpreszeros ? min(widelength(A), _sumnnzs(A, Bs...)) : widelength(A))
C = _allocres(size(A), indextypeC, entrytypeC, maxnnzC)
fpreszeros ? _map_zeropres!(f, C, A, Bs...) :
Expand All @@ -232,13 +274,20 @@ copy(bc::SpBroadcasted1) = _asksthezeros(bc) && _fillszeros(bc.f, bc.args[1]) ?
isempty(C) && return _finishempty!(C)
f = bc.f
fofnoargs = f()
if _iszero(fofnoargs) # f() is zero, so empty C
if _iszeroresult(fofnoargs) # f() is zero, so empty C
trimstorage!(C, 0)
_finishempty!(C)
else # f() is nonzero, so densify C and fill with independent calls to f()
_densestructure!(C)
storedvals(C)[1] = fofnoargs
broadcast!(f, view(storedvals(C), 2:length(storedvals(C))))
_checkdensestructure(C)
firstvalue = convert(eltype(C), fofnoargs)
try
_densestructure!(C)
storedvals(C)[1] = firstvalue
broadcast!(f, view(storedvals(C), 2:length(storedvals(C))))
catch
_is_fixed(C) || _emptypattern!(C)
rethrow()
end
end
return _checkbuffers(C)
end
Expand All @@ -254,7 +303,7 @@ function _diffshape_broadcast(f::Tf, A::SparseVecOrMat, Bs::Vararg{SparseVecOrMa
# Avoid calculating f(zero) unless necessary as it may fail.
r = if _haszeros(A) && all(_haszeros, Bs)
fofzeros = f(_zeros_eltypes(A, Bs...)...)
_iszero(fofzeros) ?
_iszeroresult(fofzeros) ?
_broadcast_zeropres!(f, _allocres(shapeC, indextypeC, entrytypeC, maxnnzC), A, Bs...) :
_broadcast_notzeropres!(f, fofzeros, _allocres(shapeC, indextypeC, entrytypeC, _densennz(shapeC)), A, Bs...)
else
Expand Down Expand Up @@ -328,12 +377,15 @@ Densifies `C`, storing `fillvalue` in place of each unstored entry in `A` and
function _map_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, A::SparseVecOrMat) where Tf
# Build dense matrix structure in C, expanding storage if necessary
_densestructure!(C)
# Below, and in the other `notzeropres` kernels, a value is left as filled only when it is
# identical to `fillvalue`: one that compares equal may be of another type or sign, and
# the comparison itself may be `missing`.
# Populate values
fill!(storedvals(C), fillvalue)
@inbounds for (j, jo) in zip(columns(C), _densecoloffsets(C))
for Ak in colrange(A, j)
Cx = f(storedvals(A)[Ak])
Cx != fillvalue && (storedvals(C)[jo + storedinds(A)[Ak]] = Cx)
Cx !== fillvalue && (storedvals(C)[jo + storedinds(A)[Ak]] = Cx)
end
end
# NOTE: Combining the fill! above into the loop above to avoid multiple sweeps over /
Expand Down Expand Up @@ -419,7 +471,7 @@ function _map_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, A::SparseVecOrMa
Cx, Ci = f(zero(eltype(A)), storedvals(B)[Bk]), Bi
Bk += oneunit(Bk); Bi = Bk < stopBk ? W(storedinds(B)[Bk]) : rowsentinelB
end
Cx != fillvalue && (storedvals(C)[jo + Ci] = Cx)
Cx !== fillvalue && (storedvals(C)[jo + Ci] = Cx)
end
end
return _checkbuffers(C)
Expand Down Expand Up @@ -475,7 +527,7 @@ function _map_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, As::Vararg{Spars
while activerow < rowsentinel
vals, ks, rows = _fusedupdate_all(rowsentinel, activerow, rows, ks, stopks, As)
Cx = f(vals...)
Cx != fillvalue && (storedvals(C)[jo + activerow] = Cx)
Cx !== fillvalue && (storedvals(C)[jo + activerow] = Cx)
activerow = min(rows...)
end
end
Expand Down Expand Up @@ -615,7 +667,7 @@ function _broadcast_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, A::SparseV
bccolrangejA = numcols(A) == 1 ? colrange(A, 1) : colrange(A, j)
for Ak in bccolrangejA
Cx, Ci = f(storedvals(A)[Ak]), storedinds(A)[Ak]
Cx != fillvalue && (storedvals(C)[jo + Ci] = Cx)
Cx !== fillvalue && (storedvals(C)[jo + Ci] = Cx)
end
end
# Cases with vertical expansion
Expand All @@ -625,7 +677,7 @@ function _broadcast_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, A::SparseV
Ak, stopAk = numcols(A) == 1 ? (colstartind(A, 1), colboundind(A, 1)) : (colstartind(A, j), colboundind(A, j))
Ax = Ak < stopAk ? svA[Ak] : zero(eltype(A))
fofAx = f(Ax)
if fofAx != fillvalue
if fofAx !== fillvalue
for i in (jo + 1):(jo + numrows(C))
svC[i] = fofAx
end
Expand Down Expand Up @@ -842,7 +894,7 @@ function _broadcast_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, A::SparseV
Ak += oneunit(Ak); Ai = Ak < stopAk ? W(storedinds(A)[Ak]) : rowsentinelA
Bk += oneunit(Bk); Bi = Bk < stopBk ? W(storedinds(B)[Bk]) : rowsentinelB
end
Cx != fillvalue && (storedvals(C)[jo + Ci] = Cx)
Cx !== fillvalue && (storedvals(C)[jo + Ci] = Cx)
end
end
# Cases with vertical expansion
Expand All @@ -853,7 +905,7 @@ function _broadcast_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, A::SparseV
Ax = Ak < stopAk ? storedvals(A)[Ak] : zero(eltype(A))
Bx = Bk < stopBk ? storedvals(B)[Bk] : zero(eltype(B))
Cx = f(Ax, Bx)
if Cx != fillvalue
if Cx !== fillvalue
for Ck::Int in (jo + 1):(jo + numrows(C))
storedvals(C)[Ck] = Cx
end
Expand All @@ -865,10 +917,10 @@ function _broadcast_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, A::SparseV
Bk, stopBk = numcols(B) == 1 ? (colstartind(B, 1), colboundind(B, 1)) : (colstartind(B, j), colboundind(B, j))
Ax = Ak < stopAk ? storedvals(A)[Ak] : zero(eltype(A))
fvAzB = f(Ax, zero(eltype(B)))
if fvAzB == fillvalue
if fvAzB === fillvalue
while Bk < stopBk
Cx = f(Ax, storedvals(B)[Bk])
Cx != fillvalue && (storedvals(C)[jo + storedinds(B)[Bk]] = Cx)
Cx !== fillvalue && (storedvals(C)[jo + storedinds(B)[Bk]] = Cx)
Bk += oneunit(Bk)
end
else
Expand All @@ -880,7 +932,7 @@ function _broadcast_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, A::SparseV
else
Cx = fvAzB
end
Cx != fillvalue && (storedvals(C)[jo + Ci] = Cx)
Cx !== fillvalue && (storedvals(C)[jo + Ci] = Cx)
end
end
end
Expand All @@ -890,10 +942,10 @@ function _broadcast_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, A::SparseV
Bk, stopBk = numcols(B) == 1 ? (colstartind(B, 1), colboundind(B, 1)) : (colstartind(B, j), colboundind(B, j))
Bx = Bk < stopBk ? storedvals(B)[Bk] : zero(eltype(B))
fzAvB = f(zero(eltype(A)), Bx)
if fzAvB == fillvalue
if fzAvB === fillvalue
while Ak < stopAk
Cx = f(storedvals(A)[Ak], Bx)
Cx != fillvalue && (storedvals(C)[jo + storedinds(A)[Ak]] = Cx)
Cx !== fillvalue && (storedvals(C)[jo + storedinds(A)[Ak]] = Cx)
Ak += oneunit(Ak)
end
else
Expand All @@ -905,7 +957,7 @@ function _broadcast_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, A::SparseV
else
Cx = fzAvB
end
Cx != fillvalue && (storedvals(C)[jo + Ci] = Cx)
Cx !== fillvalue && (storedvals(C)[jo + Ci] = Cx)
end
end
end
Expand Down Expand Up @@ -955,7 +1007,14 @@ _finishempty!(C::AbstractCompressedVector) = C
_finishempty!(C::AbstractSparseMatrixCSC) = (fill!(getcolptr(C), 1); C)

# special case - vector outer product
_copy(f::typeof(*), x::SparseVectorOrView, y::AdjOrTransSparseVectorOrView) = _outer(x, y)
_copy(f::typeof(*), x::SparseVectorOrView, y::AdjOrTransSparseVectorOrView) =
_outerkeepszeros(x, y) ? _outer(x, y) : _copy(f, _sparsifystructured(x), _sparsifystructured(y))
# `_outer` multiplies stored entries with stored entries alone, which gives the product when
# every stored entry times a zero is a zero. An `Inf` or a `NaN` times a zero is not.
function _outerkeepszeros(x, y)
zx, zy, op = _densezero(x), _densezero(y), wrapperop(y)
return all(v -> _iszero(v * zy), nonzeros(x)) && all(v -> _iszero(zx * op(v)), nonzeros(parent(y)))
end
@inline _outer(x::SparseVectorOrView, y::Adjoint) = return _outer(conj, x, parent(y))
@inline _outer(x::SparseVectorOrView, y::Transpose) = return _outer(identity, x, parent(y))
function _outer(trans::Tf, x, y) where Tf
Expand Down Expand Up @@ -1079,11 +1138,11 @@ function _broadcast_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, As::Vararg
rows = _initrowforcol_all(j, rowsentinel, isemptys, expandsverts, ks, As)
defaultCx = f(defargs...)
activerow = min(rows...)
if defaultCx == fillvalue # fillvalue-preserving column scan
if defaultCx === fillvalue # fillvalue-preserving column scan
while activerow < rowsentinel
args, ks, rows = _fusedupdatebc_all(rowsentinel, activerow, rows, defargs, ks, stopks, As)
Cx = f(args...)
Cx != fillvalue && (storedvals(C)[jo + activerow] = Cx)
Cx !== fillvalue && (storedvals(C)[jo + activerow] = Cx)
activerow = min(rows...)
end
else # fillvalue-non-preserving column scan
Expand All @@ -1095,7 +1154,7 @@ function _broadcast_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, As::Vararg
else
Cx = defaultCx
end
Cx != fillvalue && (storedvals(C)[jo + Ci] = Cx)
Cx !== fillvalue && (storedvals(C)[jo + Ci] = Cx)
end
end
end
Expand Down Expand Up @@ -1192,15 +1251,14 @@ end
Base.Broadcast.check_broadcast_axes(axes(dest), As...)
_is_fixed(dest) && return _copyto_fixed_broadcast!(f, dest, As...)
# Avoid calculating f(zero) unless necessary as it may fail.
all(_haszeros, As) || return _broadcast_zeropres!(f, dest, As...)
all(_haszeros, As) || return _guardpattern(_broadcast_zeropres!, f, dest, As...)
fofzeros = f(_zeros_eltypes(As...)...)
if _iszero(fofzeros)
return _broadcast_zeropres!(f, dest, As...)
if _iszeroresult(fofzeros)
return _guardpattern(_broadcast_zeropres!, f, dest, As...)
else
return _broadcast_notzeropres!(f, fofzeros, dest, As...)
return _guardfill(_broadcast_notzeropres!, f, fofzeros, dest, As...)
end
end

# A fixed `dest` whose pattern the kernels cannot write as they go gets the result from a
# writable array of its shape through `_copyto_fixed!`, which checks the pattern and the
# values before it writes. Every nonzero of the result must lie in the pattern of `dest`,
Expand Down Expand Up @@ -1335,11 +1393,13 @@ _fillszeros(f, args...) = false
# A function that throws at the zeros, or an element type without a zero, is left to the
# sparse kernels, which evaluate the function there only if some entry is not stored.
function _preserveszeros(f::F, args::Vararg{Any,N}) where {F,N}
local fofzeros
try
return _iszero(f(_zeros_eltypes(args...)...))
fofzeros = f(_zeros_eltypes(args...)...)
catch
return true
end
return _iszeroresult(fofzeros)
end

# The sparse kernels evaluate the function once at the zeros and then at the stored entries
Expand All @@ -1349,8 +1409,12 @@ end
function _densenonzero(bc::Broadcasted, f::F, args::Vararg{Any,N}) where {F,N}
T = Base.promote_typejoin_union(Base.promote_op(f, map(eltype, args)...))
(isconcretetype(T) && T !== Bool) || return _densebroadcast(bc)
return _storedasdense(_copy(f, map(_sparsifystructured, args)...))
As = map(_sparsifystructured, args)
# the kernels store every entry, which the index type of the arguments may not address
_addresses(_promote_indtype(As...), map(length, axes(bc))) || return _densebroadcast(bc)
return _storedasdense(_copy(f, As...))
end
_addresses(::Type{Ti}, shape) where {Ti} = !(Ti <: Base.BitInteger) || prod(Int64, shape) < typemax(Ti)
_storedasdense(C::Union{SparseVector,SparseMatrixCSC}) =
length(getnzval(C)) == nnz(C) == length(C) ? reshape(getnzval(C), size(C)) : Array(C)
_storedasdense(C) = Array(C)
Expand Down Expand Up @@ -1387,6 +1451,12 @@ _densifysparse(x) = x
_sparsifystructured(M::AbstractMatrix) = _isdenselike(M) ? _fullystored(M) : SparseMatrixCSC(M)
_sparsifystructured(V::AbstractVector) = _isdenselike(V) ? _fullystored(V) : SparseVector(V)
_sparsifystructured(S::SparseVecOrMat) = S
# stores the whole diagonal, as for the other banded types: a kernel that read a `-0.0` there
# as a structural zero would lose its sign
function _sparsifystructured(D::Diagonal)
n = size(D, 1)
return SparseMatrixCSC(n, n, collect(1:(n + 1)), collect(1:n), copyto!(Vector{eltype(D)}(undef, n), D.diag))
end
_sparsifystructured(S::SparseViewOfColumns) = copy(S)
_sparsifystructured(x) = x

Expand Down
Loading
Loading