nan propagation in matrix multiplication

[ivirshup]

Describe the bug

Matrix multiplication does not propagate nan like numpy does.

To Reproduce

import numpy as np
import sparse

A = np.eye(4)
A[2, 2] = np.nan

np.eye(4) @ A
# array([[ 1.,  0., nan,  0.],
#        [ 0.,  1., nan,  0.],
#        [ 0.,  0., nan,  0.],
#        [ 0.,  0., nan,  1.]])

(sparse.eye(4) @ sparse.COO(A)).todense()
# array([[ 1.,  0.,  0.,  0.],
#        [ 0.,  1.,  0.,  0.],
#        [ 0.,  0., nan,  0.],
#        [ 0.,  0.,  0.,  1.]])

Expected behavior

I would expect the same values from numpy and sparse. However, I see that nan breaks the expected sparsity pattern, which could make efficient implementation difficult. Maybe a warning or error would be appropriate?

System

• OS and version: macOS 10.15.3
• sparse version (sparse.__version__) '0.9.1'
• NumPy version (np.__version__) '1.18.3'
• Numba version (numba.__version__) '0.49.0'

Additional context

• Scipy sparse matrices have the same behavior: BUG: scipy.sparse.csc_matrix: matrix multiplication with nan scipy/scipy#7532
• I believe a similar topic was touched on in Support Everything that XArray Expects #1, for checking results of operations against the fill value.

[hameerabbasi]

Yes, this is because we (perhaps incorrectly) assume that 0 * x == 0 in matmul, and nowhere else.

[hameerabbasi]

Is there a use-case for this behaviour?

[ivirshup]

I'm doing outer joins, which I've ended up implementing with matrix multiplication since that works fairly consistently across array types. This was an inconsistency I came across. I've already worked around it, just thought it would be good for this behavior to be defined since I wasn't expecting this.

FWIW, Julia sparse arrays do the same thing.

using LinearAlgebra, SparseArrays

S = sparse(1.0I, 4, 4)
S[3, 3] = NaN
sparse(I, 4, 4) * S
# 4×4 SparseMatrixCSC{Float64,Int64} with 4 stored entries:
#   [1, 1]  =  1.0
#   [2, 2]  =  1.0
#   [3, 3]  =  NaN
#   [4, 4]  =  1.0

[hameerabbasi]

A short-term solution here would be to have a warning (like you said), (PRs welcome). If we want a long-term solution, we're looking for #365.