diff --git a/doc/extending/other_ops.rst b/doc/extending/other_ops.rst index 6bfc66b341..ebd450b003 100644 --- a/doc/extending/other_ops.rst +++ b/doc/extending/other_ops.rst @@ -96,7 +96,7 @@ You can extract the 4 fields with :func:`PyTensor.sparse.basic.csm_shape` to extract the individual fields. -You can look at the `AddSD` sparse `Op` for an example with C code. It implements +You can look at the `AddSparseDense` sparse `Op` for an example with C code. It implements the addition of a sparse matrix with a dense matrix. Sparse Tests diff --git a/pytensor/link/numba/dispatch/sparse/basic.py b/pytensor/link/numba/dispatch/sparse/basic.py index c43e8934d9..239aeb91f2 100644 --- a/pytensor/link/numba/dispatch/sparse/basic.py +++ b/pytensor/link/numba/dispatch/sparse/basic.py @@ -18,6 +18,7 @@ CSMProperties, DenseFromSparse, Diag, + EnsureSortedIndices, GetItem2d, GetItem2Lists, GetItem2ListsGrad, @@ -27,6 +28,7 @@ HStack, RowScaleCSC, SparseFromDense, + SumDuplicates, Transpose, VStack, ) @@ -69,6 +71,57 @@ def csm_constructor(data, indices, indptr, shape): return csm_constructor +@register_funcify_default_op_cache_key(EnsureSortedIndices) +def numba_funcify_EnsureSortedIndices(op, node, **kwargs): + inplace = op.inplace + + @numba_basic.numba_njit + def ensure_sorted_indices(x): + if inplace: + x.sort_indices() + return x + return x.sorted_indices() + + return ensure_sorted_indices + + +@register_funcify_default_op_cache_key(SumDuplicates) +def numba_funcify_SumDuplicates(op, node, **kwargs): + format = node.outputs[0].type.format + + @numba_basic.numba_njit + def sum_duplicates(x): + result = x.sorted_indices() + data = result.data + indices = result.indices.view(np.uint32) + indptr = result.indptr.view(np.uint32) + new_indptr = np.empty_like(indptr) + new_indptr[0] = 0 + nnz = 0 + for major in range(len(indptr) - 1): + for ptr in range(indptr[major], indptr[major + 1]): + index = indices[ptr] + if nnz > new_indptr[major] and indices[nnz - 1] == index: + data[nnz - 1] += data[ptr] + else: + data[nnz] = data[ptr] + indices[nnz] = index + nnz += 1 + new_indptr[major + 1] = nnz + + components = ( + data[:nnz], + indices[:nnz].view(np.int32), + new_indptr.view(np.int32), + ) + + if format == "csr": + return sp.sparse.csr_matrix(components, shape=x.shape) + return sp.sparse.csc_matrix(components, shape=x.shape) + + return sum_duplicates + + @register_funcify_default_op_cache_key(Cast) def numba_funcify_Cast(op, node, **kwargs): inp_dtype = node.inputs[0].type.dtype diff --git a/pytensor/link/numba/dispatch/sparse/math.py b/pytensor/link/numba/dispatch/sparse/math.py index 9b74e433c0..61dac80e23 100644 --- a/pytensor/link/numba/dispatch/sparse/math.py +++ b/pytensor/link/numba/dispatch/sparse/math.py @@ -11,10 +11,15 @@ register_funcify_default_op_cache_key, ) from pytensor.sparse import ( + AddSparseDense, + AddSparseSparse, + AddSparseSparseData, Dot, - SparseDenseMultiply, - SparseDenseVectorMultiply, + MultiplySparseDense, + MultiplySparseDenseVector, + MultiplySparseSparse, SpSum, + StructuredAddSparseDenseVector, StructuredDot, StructuredDotGradCSC, StructuredDotGradCSR, @@ -32,9 +37,9 @@ def perform(x): return perform -@register_funcify_default_op_cache_key(SparseDenseMultiply) -@register_funcify_default_op_cache_key(SparseDenseVectorMultiply) -def numba_funcify_SparseDenseMultiply(op, node, **kwargs): +@register_funcify_default_op_cache_key(MultiplySparseDense) +@register_funcify_default_op_cache_key(MultiplySparseDenseVector) +def numba_funcify_MultiplySparseDense(op, node, **kwargs): x, y = node.inputs [z] = node.outputs out_dtype = z.type.dtype @@ -649,3 +654,652 @@ def grad_spmspm_csc(x_indices, x_ptr, y, g_xy): return output return grad_spmspm_csc, cache_key + + +@register_funcify_and_cache_key(AddSparseSparse) +def numba_funcify_AddSparseSparse(op, node, **kwargs): + # NOTE: There are promotions to uint64 in a few places. + # The reason is numba promotes uint32 + uint32 to uint64, + # while uint64 + uint64 is just as fast. + # The benefit comes from using np.uint64(1) instead of 1 in increments: + # uint32/uint64 + 1 becomes signed, adding negative-index handling to array access, + # while uint64 + uint64 remains unsigned. + # Also, there are scalar promotions to the output type. This is more efficient that + # promoting the whole array at once, which requires additional memory allocation. + + x, y = node.inputs + [z] = node.outputs + out_dtype = z.type.dtype + out_type = np.dtype(out_dtype).type + + x_format = x.type.format + y_format = y.type.format + + # The output `z` has the same format as `x`. + if x_format == "csr": + constructor = sp.csr_matrix + else: + constructor = sp.csc_matrix + + @numba_basic.numba_njit + def add_sparse_sparse_canonical(x, y): + x_data = x.data + y_data = y.data + x_indices = x.indices.view(np.uint32) + y_indices = y.indices.view(np.uint32) + x_indptr = x.indptr.view(np.uint32) + y_indptr = y.indptr.view(np.uint32) + + n_major = len(x_indptr) - 1 + capacity = len(x_data) + len(y_data) + z_indptr = np.empty(n_major + 1, dtype=np.uint32) + z_indices = np.empty(capacity, dtype=np.uint32) + z_data = np.empty(capacity, dtype=out_dtype) + + nnz = np.uint64(0) + z_indptr[0] = 0 + for major in range(n_major): + x_ptr = np.uint64(x_indptr[major]) + x_end = np.uint64(x_indptr[major + 1]) + y_ptr = np.uint64(y_indptr[major]) + y_end = np.uint64(y_indptr[major + 1]) + + while x_ptr < x_end and y_ptr < y_end: + if x_indices[x_ptr] == y_indices[y_ptr]: + index = x_indices[x_ptr] + # Numba promotes narrow integer types. + # The outer cast ensures the result remains in the desired output dtype. + value = out_type(out_type(x_data[x_ptr]) + out_type(y_data[y_ptr])) + x_ptr += np.uint64(1) + y_ptr += np.uint64(1) + elif x_indices[x_ptr] < y_indices[y_ptr]: + index = x_indices[x_ptr] + value = out_type(x_data[x_ptr]) + x_ptr += np.uint64(1) + else: + index = y_indices[y_ptr] + value = out_type(y_data[y_ptr]) + y_ptr += np.uint64(1) + z_indices[nnz] = index + z_data[nnz] = value + nnz += np.uint64(value != 0) + + while x_ptr < x_end: + value = out_type(x_data[x_ptr]) + z_indices[nnz] = x_indices[x_ptr] + z_data[nnz] = value + nnz += np.uint64(value != 0) + x_ptr += np.uint64(1) + + while y_ptr < y_end: + value = out_type(y_data[y_ptr]) + z_indices[nnz] = y_indices[y_ptr] + z_data[nnz] = value + nnz += np.uint64(value != 0) + y_ptr += np.uint64(1) + + z_indptr[major + 1] = nnz + + components = ( + z_data[:nnz], + z_indices[:nnz].view(np.int32), + z_indptr.view(np.int32), + ) + return constructor(components, shape=x.shape) + + # Sparse indices are int32, so the top two uint32 values are free sentinels. + # This approach allows us to have unsigned `next_index` and unsigned `head` + not_visited = np.uint32(2**32 - 1) + end_of_list = np.uint32(2**32 - 2) + + @numba_basic.numba_njit + def add_sparse_sparse_general(x, y): + x_data = x.data + y_data = y.data + x_indices = x.indices.view(np.uint32) + y_indices = y.indices.view(np.uint32) + x_indptr = x.indptr.view(np.uint32) + y_indptr = y.indptr.view(np.uint32) + n_major = len(x_indptr) - 1 + + if x_format == "csr": + n_minor = x.shape[1] + else: + n_minor = x.shape[0] + + next_index = np.full(n_minor, not_visited, dtype=np.uint32) + x_row = np.empty(n_minor, dtype=out_dtype) + y_row = np.empty(n_minor, dtype=out_dtype) + capacity = len(x_data) + len(y_data) + z_data = np.empty(capacity, dtype=out_dtype) + z_indices = np.empty(capacity, dtype=np.int32) + z_indptr = np.empty(n_major + 1, dtype=np.int32) + nnz = np.uint64(0) + z_indptr[0] = 0 + + for major in range(n_major): + head = end_of_list + length = np.uint64(0) + + for ptr in range(x_indptr[major], x_indptr[major + 1]): + index = x_indices[ptr] + if next_index[index] == not_visited: + # Initialize accumulators only for indices visited in this row. + x_row[index] = out_type(x_data[ptr]) + y_row[index] = 0 + next_index[index] = head + head = index + length += np.uint64(1) + else: + x_row[index] += out_type(x_data[ptr]) + + for ptr in range(y_indptr[major], y_indptr[major + 1]): + index = y_indices[ptr] + if next_index[index] == not_visited: + x_row[index] = 0 + y_row[index] = out_type(y_data[ptr]) + next_index[index] = head + head = index + length += np.uint64(1) + else: + y_row[index] += out_type(y_data[ptr]) + + for _ in range(length): + value = out_type(x_row[head] + y_row[head]) + if value != 0: + z_indices[nnz] = head + z_data[nnz] = value + nnz += np.uint64(1) + index = head + head = next_index[index] + next_index[index] = not_visited + + z_indptr[major + 1] = nnz + return constructor((z_data[:nnz], z_indices[:nnz], z_indptr), shape=x.shape) + + cache_version = 1 + cache_key = sha256( + str((AddSparseSparse, x_format, y_format, out_dtype, cache_version)).encode() + ).hexdigest() + + @numba_basic.numba_njit + def add_sparse_sparse(x, y): + assert x.shape == y.shape + + if x_format != y_format: + # Addition is CSR-CSR or CSC-CSC + if x_format == "csr": + y = y.tocsr() + else: + y = y.tocsc() + + if x.has_canonical_format and y.has_canonical_format: + return add_sparse_sparse_canonical(x, y) + + return add_sparse_sparse_general(x, y) + + return add_sparse_sparse, cache_key + + +@register_funcify_and_cache_key(MultiplySparseSparse) +def numba_funcify_MultiplySparseSparse(op, node, **kwargs): + # NOTE: There are promotions to uint64 in a few places. + # The reason is numba promotes uint32 + uint32 to uint64, + # while uint64 + uint64 is just as fast. + # The benefit comes from using np.uint64(1) instead of 1 in increments: + # uint32/uint64 + 1 becomes signed, adding negative-index handling to array access, + # while uint64 + uint64 remains unsigned. + # Also, there are scalar promotions to the output type. This is more efficient that + # promoting the whole array at once, which requires additional memory allocation. + + x, y = node.inputs + [z] = node.outputs + out_dtype = z.type.dtype + out_type = np.dtype(out_dtype).type + + x_format = x.type.format + y_format = y.type.format + + # The output `z` has the same format as `x`. + if x_format == "csr": + constructor = sp.csr_matrix + else: + constructor = sp.csc_matrix + + @numba_basic.numba_njit + def multiply_sparse_sparse_canonical(x, y): + x_data = x.data + y_data = y.data + x_indices = x.indices.view(np.uint32) + y_indices = y.indices.view(np.uint32) + x_indptr = x.indptr.view(np.uint32) + y_indptr = y.indptr.view(np.uint32) + n_major = len(x_indptr) - 1 + capacity = len(x_data) + len(y_data) + z_indptr = np.empty(n_major + 1, dtype=np.uint32) + z_indices = np.empty(capacity, dtype=np.uint32) + z_data = np.empty(capacity, dtype=out_dtype) + + nnz = np.uint64(0) + z_indptr[0] = 0 + for major in range(n_major): + x_ptr = np.uint64(x_indptr[major]) + x_end = np.uint64(x_indptr[major + 1]) + y_ptr = np.uint64(y_indptr[major]) + y_end = np.uint64(y_indptr[major + 1]) + + while x_ptr < x_end and y_ptr < y_end: + # Numba promotes narrow integer types. + # The outer casts ensure the result remains in the desired output dtype. + if x_indices[x_ptr] == y_indices[y_ptr]: + index = x_indices[x_ptr] + value = out_type(out_type(x_data[x_ptr]) * out_type(y_data[y_ptr])) + x_ptr += np.uint64(1) + y_ptr += np.uint64(1) + elif x_indices[x_ptr] < y_indices[y_ptr]: + index = x_indices[x_ptr] + value = out_type(out_type(x_data[x_ptr]) * out_type(0)) + x_ptr += np.uint64(1) + else: + index = y_indices[y_ptr] + value = out_type(out_type(0) * out_type(y_data[y_ptr])) + y_ptr += np.uint64(1) + + if value != 0: + z_indices[nnz] = index + z_data[nnz] = value + nnz += np.uint64(1) + + while x_ptr < x_end: + value = out_type(out_type(x_data[x_ptr]) * out_type(0)) + if value != 0: + z_indices[nnz] = x_indices[x_ptr] + z_data[nnz] = value + nnz += np.uint64(1) + x_ptr += np.uint64(1) + + while y_ptr < y_end: + value = out_type(out_type(0) * out_type(y_data[y_ptr])) + if value != 0: + z_indices[nnz] = y_indices[y_ptr] + z_data[nnz] = value + nnz += np.uint64(1) + y_ptr += np.uint64(1) + + z_indptr[major + 1] = nnz + + components = ( + z_data[:nnz], + z_indices[:nnz].view(np.int32), + z_indptr.view(np.int32), + ) + + return constructor(components, shape=x.shape) + + # Sparse indices are int32, so the top two uint32 values are free sentinels. + not_visited = np.uint32(2**32 - 1) + end_of_list = np.uint32(2**32 - 2) + + @numba_basic.numba_njit + def multiply_sparse_sparse_general(x, y): + x_data = x.data + y_data = y.data + x_indices = x.indices.view(np.uint32) + y_indices = y.indices.view(np.uint32) + x_indptr = x.indptr.view(np.uint32) + y_indptr = y.indptr.view(np.uint32) + n_major = len(x_indptr) - 1 + + if x_format == "csr": + n_minor = x.shape[1] + else: + n_minor = x.shape[0] + + next_index = np.full(n_minor, not_visited, dtype=np.uint32) + x_row = np.empty(n_minor, dtype=out_dtype) + y_row = np.empty(n_minor, dtype=out_dtype) + capacity = len(x_data) + len(y_data) + z_data = np.empty(capacity, dtype=out_dtype) + z_indices = np.empty(capacity, dtype=np.int32) + z_indptr = np.empty(n_major + 1, dtype=np.int32) + + nnz = np.uint64(0) + z_indptr[0] = 0 + for major in range(n_major): + head = end_of_list + length = np.uint64(0) + + for ptr in range(x_indptr[major], x_indptr[major + 1]): + index = x_indices[ptr] + if next_index[index] == not_visited: + # Initialize accumulators only for indices visited in this row. + x_row[index] = out_type(x_data[ptr]) + y_row[index] = 0 + next_index[index] = head + head = index + length += np.uint64(1) + else: + x_row[index] += out_type(x_data[ptr]) + + for ptr in range(y_indptr[major], y_indptr[major + 1]): + index = y_indices[ptr] + if next_index[index] == not_visited: + x_row[index] = 0 + y_row[index] = out_type(y_data[ptr]) + next_index[index] = head + head = index + length += np.uint64(1) + else: + y_row[index] += out_type(y_data[ptr]) + + for _ in range(length): + value = out_type(x_row[head] * y_row[head]) + if value != 0: + z_indices[nnz] = head + z_data[nnz] = value + nnz += np.uint64(1) + index = head + head = next_index[index] + next_index[index] = not_visited + z_indptr[major + 1] = nnz + + return constructor((z_data[:nnz], z_indices[:nnz], z_indptr), shape=x.shape) + + cache_version = 1 + cache_key = sha256( + str( + (MultiplySparseSparse, x_format, y_format, out_dtype, cache_version) + ).encode() + ).hexdigest() + + @numba_basic.numba_njit + def multiply_sparse_sparse(x, y): + assert x.shape == y.shape + + if x_format != y_format: + # Multiplication is CSR-CSR or CSC-CSC + if x_format == "csr": + y = y.tocsr() + else: + y = y.tocsc() + + if x.has_canonical_format and y.has_canonical_format: + return multiply_sparse_sparse_canonical(x, y) + + return multiply_sparse_sparse_general(x, y) + + return multiply_sparse_sparse, cache_key + + +@register_funcify_default_op_cache_key(AddSparseDense) +def numba_funcify_AddSparseDense(op, node, **kwargs): + x, y = node.inputs + [z] = node.outputs + out_dtype = z.type.dtype + x_format = x.type.format + same_dtype = y.type.dtype == out_dtype + + @numba_basic.numba_njit + def add_sparse_dense(x, y): + assert x.shape == y.shape + if same_dtype: + z = y.copy() + else: + z = y.astype(out_dtype) # Also generates a copy + + M = x.shape[0] + N = x.shape[1] + data = x.data + indices = x.indices.view(np.uint32) + indptr = x.indptr.view(np.uint32) + + # Since `z` already has the output type, and we use `+=`, we don't need to cast `data`. + # Numba emits LLVM that automatically handles the cast. + if x_format == "csc": + for j in range(N): + for i_idx in range(indptr[j], indptr[j + 1]): + i = indices[i_idx] + z[i, j] += data[i_idx] + else: + for i in range(M): + for j_idx in range(indptr[i], indptr[i + 1]): + j = indices[j_idx] + z[i, j] += data[j_idx] + + return z + + return add_sparse_dense + + +@register_funcify_default_op_cache_key(AddSparseSparseData) +def numba_funcify_AddSparseSparseData(op, node, **kwargs): + if node.outputs[0].type.format == "csr": + constructor = sp.csr_matrix + else: + constructor = sp.csc_matrix + + @numba_basic.numba_njit + def add_sparse_sparse_data(x, y): + assert x.shape == y.shape + assert x.data.shape == y.data.shape + + # The Op forces x.data and y.data to be of the same type, no need to cast. + data = np.empty_like(x.data) + for ptr in range(len(data)): + data[ptr] = x.data[ptr] + y.data[ptr] + + return constructor((data, x.indices.copy(), x.indptr.copy()), shape=x.shape) + + return add_sparse_sparse_data + + +@register_funcify_and_cache_key(StructuredAddSparseDenseVector) +def numba_funcify_StructuredAddSparseDenseVector(op, node, **kwargs): + # NOTE: There are promotions to uint64 in a few places. + # The reason is numba promotes uint32 + uint32 to uint64, + # while uint64 + uint64 is just as fast. + # The benefit comes from using np.uint64(1) instead of 1 in increments: + # uint32/uint64 + 1 becomes signed, adding negative-index handling to array access, + # while uint64 + uint64 remains unsigned. + + [z] = node.outputs + out_dtype = z.type.dtype + out_type = np.dtype(out_dtype).type + complex_dtype = np.issubdtype(np.dtype(out_dtype), np.complexfloating) + out_format = z.type.format + + cache_version = 1 + cache_key = sha256( + str((type(op), out_format, out_dtype, cache_version)).encode() + ).hexdigest() + + if out_format == "csr": + constructor = sp.csr_matrix + else: + constructor = sp.csc_matrix + + sum_duplicates = numba_basic.numba_funcify( + psb.sum_duplicates, node=psb.sum_duplicates.make_node(z) + ) + + @numba_basic.numba_njit(fastmath=False) + def add_nonfinite(value, other): + # Preserve NaN propagation from 0 * NaN and 0 * inf. + return out_type(value + (value != 0) * other) + + @numba_basic.numba_njit + def finite_kernel(x, y): + assert x.shape[1] == y.shape[0] + + # Can't reasign like in `x = sum_duplicates(x)` because Numba cannot infer a common + # type when the original indices are readonly. + if x.has_canonical_format: + x_data = x.data + x_indices = x.indices.view(np.uint32) + x_indptr = x.indptr.view(np.uint32) + else: + canonical_x = sum_duplicates(x) + x_data = canonical_x.data + x_indices = canonical_x.indices.view(np.uint32) + x_indptr = canonical_x.indptr.view(np.uint32) + + n_major = len(x_indptr) - 1 + capacity = len(x_data) + z_indptr = np.empty(n_major + 1, dtype=np.uint32) + z_indices = np.empty(capacity, dtype=np.uint32) + z_data = np.empty(capacity, dtype=out_dtype) + + nnz = np.uint64(0) + z_indptr[0] = 0 + + for major in range(n_major): + start = x_indptr[major] + end = x_indptr[major + 1] + for ptr in range(start, end): + index = x_indices[ptr] + if out_format == "csr": + column = index + else: + column = major + + value = x_data[ptr] + + # Avoid a complex multiply for the structured mask. + if complex_dtype: + if value != 0: + value = out_type(value + y[column]) + if value != 0: + z_indices[nnz] = index + z_data[nnz] = value + nnz += np.uint64(1) + else: + value = out_type(value + (value != 0) * y[column]) + if value != 0: + z_indices[nnz] = index + z_data[nnz] = value + nnz += np.uint64(1) + + z_indptr[major + 1] = nnz + + components = ( + z_data[:nnz], + z_indices[:nnz].view(np.int32), + z_indptr.view(np.int32), + ) + + return constructor(components, shape=x.shape) + + @numba_basic.numba_njit + def nonfinite_kernel(x, y, nonfinite_columns): + assert x.shape[1] == y.shape[0] + + # Can't reasign like in `x = sum_duplicates(x)` because Numba cannot infer a common + # type when the original indices are readonly. + if x.has_canonical_format: + x_data = x.data + x_indices = x.indices.view(np.uint32) + x_indptr = x.indptr.view(np.uint32) + else: + canonical_x = sum_duplicates(x) + x_data = canonical_x.data + x_indices = canonical_x.indices.view(np.uint32) + x_indptr = canonical_x.indptr.view(np.uint32) + + n_major = len(x_indptr) - 1 + if out_format == "csr": + n_minor = np.uint64(x.shape[1]) + else: + n_minor = np.uint64(x.shape[0]) + + nonfinite_columns_n = len(nonfinite_columns) + capacity = len(x_data) + x.shape[0] * nonfinite_columns_n + + z_indptr = np.empty(n_major + 1, dtype=np.uint32) + z_indices = np.empty(capacity, dtype=np.uint32) + z_data = np.empty(capacity, dtype=out_dtype) + + nnz = np.uint64(0) + z_indptr[0] = 0 + for major in range(n_major): + ptr = np.uint64(x_indptr[major]) + end = np.uint64(x_indptr[major + 1]) + extra_ptr = np.uint64(0) + + if out_format == "csr": + extra_count = np.uint64(nonfinite_columns_n) + elif not np.isfinite(y[major]): + extra_count = n_minor + else: + extra_count = np.uint64(0) + + # CSC columns with finite y need no implicit entries, even in this kernel. + if out_format == "csc" and extra_count == 0: + for data_ptr in range(ptr, end): + value = x_data[data_ptr] + value = out_type(value + (value != 0) * y[major]) + if value != 0: + z_indices[nnz] = x_indices[data_ptr] + z_data[nnz] = value + nnz += np.uint64(1) + z_indptr[major + 1] = nnz + continue + + while ptr < end or extra_ptr < extra_count: + if ptr < end: + stored_index = x_indices[ptr] + else: + stored_index = n_minor + + extra_index = n_minor + if extra_ptr < extra_count: + if out_format == "csr": + extra_index = np.uint64(nonfinite_columns[extra_ptr]) + else: + extra_index = extra_ptr + + index = min(stored_index, extra_index) + value = out_type(0) + + if stored_index == index: + value = x_data[ptr] + ptr += np.uint64(1) + + if extra_index == index: + extra_ptr += np.uint64(1) + + if out_format == "csr": + column = index + else: + column = major + + value = add_nonfinite(value, y[column]) + if value != 0: + z_indices[nnz] = index + z_data[nnz] = value + nnz += np.uint64(1) + + z_indptr[major + 1] = nnz + + components = ( + z_data[:nnz], + z_indices[:nnz].view(np.int32), + z_indptr.view(np.int32), + ) + + return constructor(components, shape=x.shape) + + @numba_basic.numba_njit + def structured_add_sparse_dense_vector(x, y): + all_finite = True + + # This is faster in numba, it's SIMD-friendly. + for y_i in y: + all_finite &= np.isfinite(y_i) + + if all_finite: + return finite_kernel(x, y) + + nonfinite_columns = np.flatnonzero(~np.isfinite(y)) + return nonfinite_kernel(x, y, nonfinite_columns) + + return structured_add_sparse_dense_vector, cache_key diff --git a/pytensor/link/numba/dispatch/sparse/variable.py b/pytensor/link/numba/dispatch/sparse/variable.py index f3c5696caa..b6069997c9 100644 --- a/pytensor/link/numba/dispatch/sparse/variable.py +++ b/pytensor/link/numba/dispatch/sparse/variable.py @@ -375,6 +375,51 @@ def copy(matrix): return copy +@overload_attribute(CSMatrixType, "has_canonical_format") +def overload_sparse_has_canonical_format(matrix): + def has_canonical_format(matrix): + indices = matrix.indices.view(np.uint32) + indptr = matrix.indptr.view(np.uint32) + for major in range(len(indptr) - 1): + start = indptr[major] + end = indptr[major + 1] + for ptr in range(start + np.uint64(1), end): + if indices[ptr - np.uint64(1)] >= indices[ptr]: + return False + return True + + return has_canonical_format + + +@overload_method(CSMatrixType, "sort_indices") +def overload_sparse_sort_indices(matrix): + def sort_indices(matrix): + data = matrix.data + indices = matrix.indices.view(np.uint32) + indptr = matrix.indptr.view(np.uint32) + for major in range(len(indptr) - 1): + start, end = indptr[major], indptr[major + 1] + # Despite branching, checking first is faster than always calling argsort. + for ptr in range(start + 1, end): + if indices[ptr - 1] > indices[ptr]: + order = np.argsort(indices[start:end]) + data[start:end] = data[start:end][order] + indices[start:end] = indices[start:end][order] + break + + return sort_indices + + +@overload_method(CSMatrixType, "sorted_indices") +def overload_sparse_sorted_indices(matrix): + def sorted_indices(matrix): + result = matrix.copy() + result.sort_indices() + return result + + return sorted_indices + + @overload_method(CSMatrixType, "astype") def overload_sparse_astype(matrix, dtype): match matrix: diff --git a/pytensor/sparse/basic.py b/pytensor/sparse/basic.py index 708ad690d9..592fe8fab6 100644 --- a/pytensor/sparse/basic.py +++ b/pytensor/sparse/basic.py @@ -1483,7 +1483,7 @@ class EnsureSortedIndices(Op): def __init__(self, inplace): self.inplace = inplace if self.inplace: - self.view_map = {0: [0]} + self.destroy_map = {0: [0]} def make_node(self, x): """ @@ -1501,7 +1501,8 @@ def perform(self, node, inputs, outputs): (x,) = inputs (z,) = outputs if self.inplace: - z[0] = x.sort_indices() + x.sort_indices() + z[0] = x else: z[0] = x.sorted_indices() @@ -1521,6 +1522,66 @@ def __str__(self): ensure_sorted_indices = EnsureSortedIndices(inplace=False) +class HasCanonicalFormat(Op): + """Read whether a sparse matrix has sorted indices and no duplicates. + + Uses SciPy's cached property without modifying the sparse components. + """ + + __props__ = () + + def make_node(self, x): + x = as_sparse_variable(x) + return Apply(self, [x], [scalar(dtype="bool")]) + + def perform(self, node, inputs, outputs): + (x,) = inputs + (z,) = outputs + z[0] = np.asarray(x.has_canonical_format, dtype="bool") + + def connection_pattern(self, node): + return [[False]] + + def pullback(self, inputs, outputs, output_grad): + return [disconnected_type()] + + def infer_shape(self, node, shapes): + return [()] + + +has_canonical_format = HasCanonicalFormat() + + +class SumDuplicates(Op): + """Sort sparse indices and sum duplicate entries without removing explicit zeros. + + The input matrix is not modified. The gradient is regular, not structured. + """ + + __props__ = () + + def make_node(self, x): + x = as_sparse_variable(x) + assert x.format in ("csr", "csc") + return Apply(self, [x], [x.type()]) + + def perform(self, node, inputs, outputs): + (x,) = inputs + (z,) = outputs + result = x.copy() + result.sum_duplicates() + z[0] = result + + def pullback(self, inputs, outputs, output_grad): + return [output_grad[0]] + + def infer_shape(self, node, shapes): + return shapes + + +sum_duplicates = SumDuplicates() + + def clean(x): """ Remove explicit zeros from a sparse matrix, and re-sort indices. diff --git a/pytensor/sparse/math.py b/pytensor/sparse/math.py index e1f986e826..dd1b35ef43 100644 --- a/pytensor/sparse/math.py +++ b/pytensor/sparse/math.py @@ -376,7 +376,7 @@ def sp_sum(x, axis=None, sparse_grad=False): return SpSum(axis, sparse_grad)(x) -class AddSS(Op): +class AddSparseSparse(Op): # add(sparse, sparse). # see the doc of add() for more detail. __props__ = () @@ -410,10 +410,10 @@ def infer_shape(self, node, shapes): return [shapes[0]] -add_s_s = AddSS() +add_s_s = AddSparseSparse() -class AddSSData(Op): +class AddSparseSparseData(Op): """Add two sparse matrices assuming they have the same sparsity pattern. Notes @@ -471,10 +471,10 @@ def infer_shape(self, node, ins_shapes): return [ins_shapes[0]] -add_s_s_data = AddSSData() +add_s_s_data = AddSparseSparseData() -class AddSD(Op): +class AddSparseDense(Op): # add(sparse, sparse). # see the doc of add() for more detail. __props__ = () @@ -513,10 +513,10 @@ def infer_shape(self, node, shapes): return [shapes[1]] -add_s_d = AddSD() +add_s_d = AddSparseDense() -class StructuredAddSV(Op): +class StructuredAddSparseDenseVector(Op): """Structured addition of a sparse matrix and a dense vector. The elements of the vector are only added to the corresponding @@ -573,7 +573,7 @@ def infer_shape(self, node, ins_shapes): return [ins_shapes[0]] -structured_add_s_v = StructuredAddSV() +structured_add_s_v = StructuredAddSparseDenseVector() def add(x, y): @@ -668,7 +668,7 @@ def sub(x, y): sub.__doc__ = subtract.__doc__ -class SparseSparseMultiply(Op): +class MultiplySparseSparse(Op): # mul(sparse, sparse) # See the doc of mul() for more detail __props__ = () @@ -703,10 +703,10 @@ def infer_shape(self, node, shapes): return [shapes[0]] -mul_s_s = SparseSparseMultiply() +mul_s_s = MultiplySparseSparse() -class SparseDenseMultiply(Op): +class MultiplySparseDense(Op): # mul(sparse, dense) # See the doc of mul() for more detail __props__ = () @@ -799,10 +799,10 @@ def infer_shape(self, node, shapes): return [shapes[0]] -mul_s_d = SparseDenseMultiply() +mul_s_d = MultiplySparseDense() -class SparseDenseVectorMultiply(Op): +class MultiplySparseDenseVector(Op): """Element-wise multiplication of sparse matrix by a broadcasted dense vector element wise. Notes @@ -811,7 +811,7 @@ class SparseDenseVectorMultiply(Op): """ - # TODO: Merge with the SparseDenseMultiply Op + # TODO: Merge with the MultiplySparseDense Op __props__ = () @@ -882,7 +882,7 @@ def infer_shape(self, node, ins_shapes): return [ins_shapes[0]] -mul_s_v = SparseDenseVectorMultiply() +mul_s_v = MultiplySparseDenseVector() def multiply(x, y): @@ -950,7 +950,7 @@ def mul(x, y): mul.__doc__ = multiply.__doc__ -class __ComparisonOpSS(Op): +class __ComparisonOpSparseSparse(Op): """ Used as a superclass for all comparisons between two sparses matrices. @@ -1000,7 +1000,7 @@ def infer_shape(self, node, ins_shapes): return [ins_shapes[0]] -class __ComparisonOpSD(Op): +class __ComparisonOpSparseDense(Op): """ Used as a superclass for all comparisons between sparse and dense matrix. @@ -1096,100 +1096,100 @@ def helper(x, y): return helper -class EqualSS(__ComparisonOpSS): +class EqualSparseSparse(__ComparisonOpSparseSparse): def comparison(self, x, y): return x == y -equal_s_s = EqualSS() +equal_s_s = EqualSparseSparse() -class EqualSD(__ComparisonOpSD): +class EqualSparseDense(__ComparisonOpSparseDense): def comparison(self, x, y): return x == y -equal_s_d = EqualSD() +equal_s_d = EqualSparseDense() -class NotEqualSS(__ComparisonOpSS): +class NotEqualSparseSparse(__ComparisonOpSparseSparse): def comparison(self, x, y): return x != y -not_equal_s_s = NotEqualSS() +not_equal_s_s = NotEqualSparseSparse() -class NotEqualSD(__ComparisonOpSD): +class NotEqualSparseDense(__ComparisonOpSparseDense): def comparison(self, x, y): return x != y -not_equal_s_d = NotEqualSD() +not_equal_s_d = NotEqualSparseDense() -class LessThanSS(__ComparisonOpSS): +class LessThanSparseSparse(__ComparisonOpSparseSparse): def comparison(self, x, y): return x < y -less_than_s_s = LessThanSS() +less_than_s_s = LessThanSparseSparse() -class LessThanSD(__ComparisonOpSD): +class LessThanSparseDense(__ComparisonOpSparseDense): def comparison(self, x, y): return x < y -less_than_s_d = LessThanSD() +less_than_s_d = LessThanSparseDense() -class GreaterThanSS(__ComparisonOpSS): +class GreaterThanSparseSparse(__ComparisonOpSparseSparse): def comparison(self, x, y): return x > y -greater_than_s_s = GreaterThanSS() +greater_than_s_s = GreaterThanSparseSparse() -class GreaterThanSD(__ComparisonOpSD): +class GreaterThanSparseDense(__ComparisonOpSparseDense): def comparison(self, x, y): return x > y -greater_than_s_d = GreaterThanSD() +greater_than_s_d = GreaterThanSparseDense() -class LessEqualSS(__ComparisonOpSS): +class LessEqualSparseSparse(__ComparisonOpSparseSparse): def comparison(self, x, y): return x <= y -less_equal_s_s = LessEqualSS() +less_equal_s_s = LessEqualSparseSparse() -class LessEqualSD(__ComparisonOpSD): +class LessEqualSparseDense(__ComparisonOpSparseDense): def comparison(self, x, y): return x <= y -less_equal_s_d = LessEqualSD() +less_equal_s_d = LessEqualSparseDense() -class GreaterEqualSS(__ComparisonOpSS): +class GreaterEqualSparseSparse(__ComparisonOpSparseSparse): def comparison(self, x, y): return x >= y -greater_equal_s_s = GreaterEqualSS() +greater_equal_s_s = GreaterEqualSparseSparse() -class GreaterEqualSD(__ComparisonOpSD): +class GreaterEqualSparseDense(__ComparisonOpSparseDense): def comparison(self, x, y): return x >= y -greater_equal_s_d = GreaterEqualSD() +greater_equal_s_d = GreaterEqualSparseDense() eq = __ComparisonSwitch(equal_s_s, equal_s_d, equal_s_d) @@ -2147,5 +2147,5 @@ def _vectorize_sparse_unsupported(op, node, *batched_inputs): ) -for _op_cls in (AddSD, SparseDenseMultiply): +for _op_cls in (AddSparseDense, MultiplySparseDense): _vectorize_node.register(_op_cls)(_vectorize_sparse_unsupported) diff --git a/pytensor/sparse/rewriting.py b/pytensor/sparse/rewriting.py index 1a2c75a6a8..d69dd4e2f1 100644 --- a/pytensor/sparse/rewriting.py +++ b/pytensor/sparse/rewriting.py @@ -66,7 +66,7 @@ def local_inplace_remove0(fgraph, node): ) -class AddSD_ccode(_NoPythonCOp): +class AddSparseDenseC(_NoPythonCOp): """ Add a sparse and a dense matrix. @@ -183,13 +183,13 @@ def c_code_cache_version(self): return (3,) -@node_rewriter([spm.AddSD]) +@node_rewriter([spm.AddSparseDense]) def local_inplace_addsd_ccode(fgraph, node): - """Rewrite to insert inplace versions of `AddSD`.""" + """Rewrite to insert inplace versions of `AddSparseDense`.""" out_dtype = ps.upcast(*[inp.type.dtype for inp in node.inputs]) if out_dtype != node.inputs[1].dtype: return - new_node = AddSD_ccode(format=node.inputs[0].type.format, inplace=True)( + new_node = AddSparseDenseC(format=node.inputs[0].type.format, inplace=True)( *node.inputs ) return [new_node] @@ -216,13 +216,13 @@ def local_dense_from_sparse_sparse_from_dense(fgraph, node): return inp.owner.inputs -@node_rewriter([spm.AddSD]) +@node_rewriter([spm.AddSparseDense]) def local_addsd_ccode(fgraph, node): """ - Convert AddSD to faster AddSD_ccode. + Convert AddSparseDense to faster AddSparseDenseC. """ - new_node = AddSD_ccode(format=node.inputs[0].type.format)(*node.inputs) + new_node = AddSparseDenseC(format=node.inputs[0].type.format)(*node.inputs) return [new_node] @@ -1595,7 +1595,7 @@ def local_mul_s_v(fgraph, node): register_specialize(local_mul_s_v, "cxx_only") -class StructuredAddSVCSR(_NoPythonCOp): +class StructuredAddSparseDenseVectorCSR(_NoPythonCOp): """Structured addition of a sparse matrix and a dense vector. The elements of the vector are are only added to the corresponding @@ -1605,7 +1605,7 @@ class StructuredAddSVCSR(_NoPythonCOp): Notes ----- - This `Op` is used as a rewritten form of `StructuredAddSV`. + This `Op` is used as a rewritten form of `StructuredAddSparseDenseVector`. """ @@ -1737,7 +1737,7 @@ def __str__(self): return self.__class__.__name__ -structured_add_s_v_csr = StructuredAddSVCSR() +structured_add_s_v_csr = StructuredAddSparseDenseVectorCSR() # register a specialization to replace diff --git a/pytensor/sparse/variable.py b/pytensor/sparse/variable.py index 7fbe9eb366..35f212cfd5 100644 --- a/pytensor/sparse/variable.py +++ b/pytensor/sparse/variable.py @@ -11,6 +11,7 @@ get_item_2lists, get_item_list, get_item_scalar, + has_canonical_format, sp_ones_like, sp_zeros_like, transpose, @@ -139,6 +140,11 @@ def shape(self): ndim = property(lambda self: self.type.ndim) dtype = property(lambda self: self.type.dtype) + @property + def has_canonical_format(self): + """Indicates if the matrix has sorted indices without duplicates.""" + return has_canonical_format(self) + # Note that the `size` attribute of sparse matrices behaves differently # from dense matrices: it is the number of elements stored in the matrix # rather than the total number of elements that may be stored. Note also diff --git a/tests/link/numba/sparse/test_basic.py b/tests/link/numba/sparse/test_basic.py index d9b167223a..d67c617ddd 100644 --- a/tests/link/numba/sparse/test_basic.py +++ b/tests/link/numba/sparse/test_basic.py @@ -80,20 +80,21 @@ def create_csr_matrix(data, indices, ind_ptr): x_data = x.data x_indptr = x.indptr - assert getrefcount(x_data) == 3 - assert getrefcount(x_indptr) == 3 + # Python 3.14+ can use borrowed references, changing the initial refcount. + data_refcount = getrefcount(x_data) + indptr_refcount = getrefcount(x_indptr) for i in range(5): a = create_csr_matrix(x.data, x.indices, x.indptr) # a.data is a view of the underlying data under x.data, but doesn't reference it directly - assert getrefcount(x_data) == 3 + assert getrefcount(x_data) == data_refcount # x_indptr is reused directly - assert getrefcount(x_indptr) == 4 + assert getrefcount(x_indptr) == indptr_refcount + 1 del a - assert getrefcount(x_data) == 3 - assert getrefcount(x_indptr) == 3 + assert getrefcount(x_data) == data_refcount + assert getrefcount(x_indptr) == indptr_refcount def test_sparse_passthrough_refcount(): @@ -104,12 +105,13 @@ def identity(a): x = sp.sparse.random(5, 5, density=0.5, format="csr") x_data = x.data - assert getrefcount(x_data) == 3 + # Python 3.14+ can use borrowed references, changing the initial refcount. + data_refcount = getrefcount(x_data) for i in range(5): identity(x) - assert getrefcount(x_data) == 3 + assert getrefcount(x_data) == data_refcount def test_sparse_shape(): @@ -676,3 +678,79 @@ def test_sparse_diag_not_square_raises(format): with pytest.raises(ValueError, match="Diag only apply on square matrix"): fn(x_test) + + +@pytest.mark.parametrize("format", ("csr", "csc")) +@pytest.mark.parametrize("dtype", ("float64", "int32", "complex64")) +@pytest.mark.parametrize("inplace", (False, True)) +def test_sparse_sorted_indices(format, dtype, inplace): + x = ps.matrix(format=format, name="x", dtype=dtype) + z = ps.EnsureSortedIndices(inplace=inplace)(x) + rng = np.random.default_rng(47) + x_test = sp.sparse.random( + 8, + 6, + density=0.5, + format=format, + dtype=dtype, + random_state=rng, + data_rvs=lambda size: rng.integers(1, 10, size=size), + ) + for start, end in zip(x_test.indptr[:-1], x_test.indptr[1:], strict=True): + x_test.indices[start:end] = x_test.indices[start:end][::-1] + x_test.data[start:end] = x_test.data[start:end][::-1] + x_test.has_sorted_indices = False + + fn, result = compare_numba_and_py_sparse([x], z, [x_test], inplace=inplace) + assert result.has_sorted_indices + assert result.nnz == x_test.nnz + + for value in [x_test.sorted_indices(), x_test[:0, :].copy(), x_test[:, :0].copy()]: + expected = value.sorted_indices() + result = fn(value) + assert result.has_sorted_indices + sparse_assert_fn(result, expected) + assert ps.SparseTensorType.may_share_memory(result, value) == inplace + + +@pytest.mark.parametrize("format", ("csr", "csc")) +@pytest.mark.parametrize("dtype", ("float64", "int32", "complex64")) +def test_sparse_sum_duplicates(format, dtype): + x = ps.matrix(format=format, name="x", dtype=dtype) + z = ps.sum_duplicates(x) + rng = np.random.default_rng(51) + x_test = sp.sparse.random( + 8, + 6, + density=0.5, + format=format, + dtype=dtype, + random_state=rng, + data_rvs=lambda size: rng.integers(1, 10, size=size), + ) + original_nnz = x_test.nnz + x_test.data = np.repeat(x_test.data, 2) + x_test.indices = np.repeat(x_test.indices, 2) + x_test.indptr *= 2 + x_test.data[:2] = [1, -1] + x_test.data[2:4] = 0 + for start, end in zip(x_test.indptr[:-1], x_test.indptr[1:], strict=True): + x_test.indices[start:end] = x_test.indices[start:end][::-1] + x_test.data[start:end] = x_test.data[start:end][::-1] + x_test.has_sorted_indices = False + x_test.has_canonical_format = False + + fn, result = compare_numba_and_py_sparse([x], z, [x_test]) + assert result.has_canonical_format + assert result.nnz == original_nnz + assert np.count_nonzero(result.data == 0) == 2 + + canonical = x_test.copy() + canonical.sum_duplicates() + for value in [canonical, x_test[:0, :].copy(), x_test[:, :0].copy()]: + expected = value.copy() + expected.sum_duplicates() + result = fn(value) + assert result.has_canonical_format + sparse_assert_fn(result, expected) + assert not ps.SparseTensorType.may_share_memory(result, value) diff --git a/tests/link/numba/sparse/test_math.py b/tests/link/numba/sparse/test_math.py index 3f9bbceff6..ae6bd56cb7 100644 --- a/tests/link/numba/sparse/test_math.py +++ b/tests/link/numba/sparse/test_math.py @@ -2,9 +2,14 @@ import pytest import scipy +import pytensor import pytensor.sparse as ps import pytensor.tensor as pt -from tests.link.numba.sparse.test_basic import compare_numba_and_py_sparse +from tests.link.numba.sparse.test_basic import ( + compare_numba_and_py_sparse, + sparse_assert_fn, +) +from tests.link.numba.test_basic import numba_mode pytestmark = pytest.mark.filterwarnings("error") @@ -165,3 +170,387 @@ def test_sparse_sum(format, axis): x_test = scipy.sparse.random(7, 5, density=0.4, format=format) compare_numba_and_py_sparse([x], z, [x_test]) + + +class TestSparseSparseBinary: + @pytest.fixture( + scope="class", + params=[ + (x_format, y_format, x_dtype, y_dtype) + for x_format in ("csr", "csc") + for y_format in ("csr", "csc") + for x_dtype, y_dtype in ( + ("float64", "float64"), + ("int64", "float32"), + ("float32", "complex64"), + ("int8", "int8"), + ("int8", "float64"), + ) + ], + ids=lambda params: "-".join(params), + ) + @staticmethod + def inputs(request): + x_format, y_format, x_dtype, y_dtype = request.param + return ( + ps.matrix(x_format, name="x", dtype=x_dtype), + ps.matrix(y_format, name="y", dtype=y_dtype), + ) + + @pytest.fixture(scope="class") + @staticmethod + def add_fn(inputs): + x, y = inputs + return pytensor.function(inputs, ps.add(x, y), mode=numba_mode) + + @pytest.fixture(scope="class") + @staticmethod + def multiply_fn(inputs): + x, y = inputs + return pytensor.function(inputs, ps.multiply(x, y), mode=numba_mode) + + @pytest.fixture + @staticmethod + def canonical_values(inputs): + x, y = inputs + rng = np.random.default_rng(47) + x_value = scipy.sparse.random( + 31, + 47, + density=0.3, + format=x.type.format, + random_state=rng, + data_rvs=lambda size: rng.integers(1, 10, size=size), + ).astype(x.dtype) + + y_value = scipy.sparse.random( + 31, + 47, + density=0.5, + format=y.type.format, + random_state=rng, + data_rvs=lambda size: rng.integers(1, 10, size=size), + ).astype(y.dtype) + + if np.issubdtype(y_value.dtype, np.inexact): + y_value.data[:] = rng.normal(size=y_value.nnz) + + if np.issubdtype(y_value.dtype, np.complexfloating): + y_value.data.imag[:] = rng.normal(size=y_value.nnz) + + return x_value, y_value + + @pytest.fixture( + params=[ + pytest.param((([2, 3, 4], [2, 0, 1], [0, 3, 3]), None), id="unsorted"), + pytest.param((([2, 3, 4], [0, 0, 2], [0, 3, 3]), None), id="duplicates"), + pytest.param((([2, 0, 3], [0, 1, 2], [0, 3, 3]), None), id="explicit-zero"), + pytest.param((([2, 3], [0, 2], [0, 2, 2]), None), id="cancellation"), + pytest.param((([64, 1, 0], [0, 2, 1], [0, 2, 3]), None), id="overflow"), + pytest.param( + (([100, 100], [0, 0], [0, 2, 2]), None), id="duplicate-overflow" + ), + pytest.param( + ( + ([2, 3, -2, 0, 4], [2, 0, 2, 1, 1], [0, 4, 5]), + ([1, -3, 2, 5], [2, 2, 0, 0], [0, 2, 4]), + ), + id="distinct-patterns", + ), + ] + ) + def special_values(self, request, inputs): + x, y = inputs + if x.type.format == "csr": + constructor = scipy.sparse.csr_matrix + shape = (2, 3) + else: + constructor = scipy.sparse.csc_matrix + shape = (3, 2) + + x_components, y_components = request.param + data, indices, indptr = x_components + x_value = constructor( + ( + np.array(data, dtype=x.dtype), + np.array(indices, dtype="int32"), + np.array(indptr, dtype="int32"), + ), + shape=shape, + ) + + if y_components is None: + y_value = (-x_value).astype(y.dtype).asformat(y.type.format) + else: + data, indices, indptr = y_components + y_value = constructor( + ( + np.array(data, dtype=y.dtype), + np.array(indices, dtype="int32"), + np.array(indptr, dtype="int32"), + ), + shape=shape, + ).asformat(y.type.format) + + return x_value, y_value + + @pytest.fixture + def empty_values(self, inputs): + x, y = inputs + shape = (0, 3) if x.type.format == "csr" else (3, 0) + return ( + scipy.sparse.csr_matrix(shape, dtype=x.dtype).asformat(x.type.format), + scipy.sparse.csr_matrix(shape, dtype=y.dtype).asformat(y.type.format), + ) + + def test_add_canonical(self, add_fn, canonical_values): + x, y = canonical_values + result = add_fn(x, y) + expected = x + y + sparse_assert_fn(result, expected) + if (x.dtype, y.dtype) == ("int64", "float32"): + np.testing.assert_allclose( + result.toarray(), expected.toarray(), rtol=1e-14, atol=0, strict=True + ) + + def test_multiply_canonical(self, multiply_fn, canonical_values): + x, y = canonical_values + result = multiply_fn(x, y) + expected = x.multiply(y) + sparse_assert_fn(result, expected) + if (x.dtype, y.dtype) == ("int64", "float32"): + np.testing.assert_allclose( + result.toarray(), expected.toarray(), rtol=1e-14, atol=0, strict=True + ) + + def test_add_special_values(self, add_fn, special_values): + x, y = special_values + result = add_fn(x, y) + # SciPy may normalize its operands in place. + sparse_assert_fn(result, x.copy() + y.copy()) + + def test_multiply_special_values(self, multiply_fn, special_values): + x, y = special_values + result = multiply_fn(x, y) + sparse_assert_fn(result, x.copy().multiply(y.copy())) + + def test_add_empty(self, add_fn, empty_values): + x, y = empty_values + sparse_assert_fn(add_fn(x, y), x + y) + + def test_multiply_empty(self, multiply_fn, empty_values): + x, y = empty_values + sparse_assert_fn(multiply_fn(x, y), x.multiply(y)) + + def test_add_preserves_inputs(self, add_fn, special_values): + x, y = special_values + x_before, y_before = x.copy(), y.copy() + add_fn(x, y) + for value, before in ((x, x_before), (y, y_before)): + np.testing.assert_array_equal(value.data, before.data) + np.testing.assert_array_equal(value.indices, before.indices) + np.testing.assert_array_equal(value.indptr, before.indptr) + + def test_multiply_preserves_inputs(self, multiply_fn, special_values): + x, y = special_values + x_before, y_before = x.copy(), y.copy() + multiply_fn(x, y) + for value, before in ((x, x_before), (y, y_before)): + np.testing.assert_array_equal(value.data, before.data) + np.testing.assert_array_equal(value.indices, before.indices) + np.testing.assert_array_equal(value.indptr, before.indptr) + + @pytest.mark.parametrize( + "format, constructor, shape", + [ + ("csr", scipy.sparse.csr_matrix, (2, 3)), + ("csc", scipy.sparse.csc_matrix, (3, 2)), + ], + ) + @pytest.mark.parametrize("indices", [[0, 2], [2, 0]]) + @np.errstate(invalid="ignore") + def test_multiply_nonfinite(self, format, constructor, shape, indices): + x = ps.matrix(format, dtype="float64") + y = ps.matrix(format, dtype="float64") + value = constructor( + ( + np.array([np.nan, np.inf]), + np.array(indices, dtype="int32"), + np.array([0, 2, 2], dtype="int32"), + ), + shape=shape, + ) + fn = compare_numba_and_py_sparse( + [x, y], + [ps.multiply(x, y), ps.multiply(y, x)], + [value, constructor(shape, dtype="float64")], + )[0] + other = constructor( + ( + np.array([np.nan]), + np.array([1], dtype="int32"), + np.array([0, 1, 1], dtype="int32"), + ), + shape=shape, + ) + actual = fn(value, other) + actual_xy = actual[0] + actual_yx = actual[1] + sparse_assert_fn(actual_xy, value.multiply(other)) + sparse_assert_fn(actual_yx, other.multiply(value)) + + +@pytest.mark.parametrize("format", ["csr", "csc"]) +@pytest.mark.parametrize("dtype", ["float32", "float64", "int8", "int32"]) +def test_sparse_sparse_data_add(format, dtype): + x = ps.matrix(format, name="x", dtype=dtype, shape=(11, 17)) + y = ps.matrix(format, name="y", dtype=dtype, shape=(11, 17)) + z = ps.add_s_s_data(x, y) + rng = np.random.default_rng(51) + integer_dtype = np.issubdtype(np.dtype(dtype), np.integer) + + def data_rvs(size): + if integer_dtype: + return rng.integers(-100, 101, size=size) + return rng.normal(size=size) + + x_test = scipy.sparse.random( + 11, + 17, + density=0.3, + format=format, + dtype=dtype, + random_state=rng, + data_rvs=data_rvs, + ) + # x and y have to follow the same sparsity pattern + y_test = x_test.copy() + y_test.data = data_rvs(size=y_test.nnz).astype(dtype) + + if dtype == "int8": + x_test.data[0] = y_test.data[0] = 100 + with np.errstate(over="ignore"): + # The Python-only execution warns about the deliberate int8 overflow + compare_numba_and_py_sparse([x, y], z, [x_test, y_test]) + else: + compare_numba_and_py_sparse([x, y], z, [x_test, y_test]) + + +@pytest.mark.parametrize("format", ["csr", "csc"]) +@pytest.mark.parametrize( + "x_dtype, y_dtype", [("float64", "float64"), ("int64", "float32")] +) +def test_sparse_dense_add(format, x_dtype, y_dtype): + x = ps.matrix(format, name="x", dtype=x_dtype, shape=(7, 9)) + y = pt.matrix("y", dtype=y_dtype, shape=(7, 9)) + outputs = [ps.add(x, y), ps.add(y, x)] + rng = np.random.default_rng(71) + x_test = scipy.sparse.random( + 7, + 9, + density=0.5, + format=format, + random_state=rng, + data_rvs=lambda size: rng.integers(1, 10, size=size), + ).astype(x_dtype) + + y_test = rng.normal(size=(7, 9)).astype(y_dtype) + compare_numba_and_py_sparse([x, y], outputs, [x_test, y_test]) + + +@pytest.mark.parametrize( + "format, constructor, shape", + [ + ("csr", scipy.sparse.csr_matrix, (2, 3)), + ("csc", scipy.sparse.csc_matrix, (3, 2)), + ], +) +def test_structured_sparse_vector_add(format, constructor, shape): + x = ps.matrix(format, name="x") + y = pt.vector("y") + z = ps.structured_add_s_v(x, y) + rng = np.random.default_rng(87) + x_test = scipy.sparse.random(11, 13, density=0.4, format=format, random_state=rng) + y_test = rng.normal(size=13) + fn = compare_numba_and_py_sparse([x, y], z, [x_test, y_test])[0] + + x_test = constructor( + ( + np.array([2.0, 3.0, -2.0, 0.0, 4.0]), + np.array([2, 0, 2, 1, 1]), + np.array([0, 4, 5]), + ), + shape=shape, + ) + y_test = np.array([-3.0, -4.0, 7.0][: shape[1]]) + original = x_test.copy() + expected = constructor(x_test + (x_test.toarray() != 0) * y_test) + actual = fn(x_test, y_test) + sparse_assert_fn(actual, expected) + np.testing.assert_array_equal(x_test.data, original.data) + np.testing.assert_array_equal(x_test.indices, original.indices) + np.testing.assert_array_equal(x_test.indptr, original.indptr) + + # Gradient graphs can embed readonly sparse indices. + x_test.indices.flags.writeable = False + x_test.indptr.flags.writeable = False + sparse_assert_fn(fn(x_test, y_test), expected) + np.testing.assert_array_equal(x_test.data, original.data) + np.testing.assert_array_equal(x_test.indices, original.indices) + np.testing.assert_array_equal(x_test.indptr, original.indptr) + + empty = constructor(shape, dtype=x.dtype) + sparse_assert_fn(fn(empty, y_test), empty) + + +@pytest.mark.parametrize("format", ["csr", "csc"]) +@pytest.mark.parametrize("dtype", ["float64", "complex128"]) +def test_structured_sparse_vector_add_nonfinite(format, dtype): + if format == "csr": + constructor = scipy.sparse.csr_matrix + else: + constructor = scipy.sparse.csc_matrix + + x = ps.matrix(format, dtype=dtype) + y = pt.vector(dtype=dtype) + x_test = scipy.sparse.csr_matrix( + ( + np.array([1, 0, -2, 0, 3, 4], dtype=dtype), + np.array([0, 1, 2, 0, 1, 3], dtype="int32"), + np.array([0, 3, 6, 6], dtype="int32"), + ), + shape=(3, 4), + ).asformat(format) + y_test = np.array([np.nan, np.inf, -np.inf, 2], dtype=dtype) + # The structured mask deliberately evaluates 0 * inf and 0 * NaN. + with np.errstate(invalid="ignore"): + comparison = compare_numba_and_py_sparse( + [x, y], ps.structured_add_s_v(x, y), [x_test, y_test] + ) + fn = comparison[0] + actual = comparison[1] + assert np.isnan(actual.toarray()[:, 0]).all() + assert np.isnan(actual.toarray()[2, :3]).all() + + for shape in [(3, 4), (0, 4), (3, 0)]: + value = constructor(shape, dtype=dtype) + vector = y_test[: shape[1]] + dense_value = value.toarray() + with np.errstate(invalid="ignore"): + dense_expected = dense_value + (dense_value != 0) * vector + + expected = constructor(dense_expected) + sparse_assert_fn(fn(value, vector), expected) + + x_test.data[:3] = [np.nan, np.inf, -np.inf] + dense_x = x_test.toarray() + with np.errstate(invalid="ignore"): + dense_expected = dense_x + (dense_x != 0) * y_test + expected = constructor(dense_expected) + sparse_assert_fn(fn(x_test, y_test), expected) + + y_test = np.arange(1, 5, dtype=dtype) + with np.errstate(invalid="ignore"): + dense_expected = dense_x + (dense_x != 0) * y_test + expected = constructor(dense_expected) + sparse_assert_fn(fn(x_test, y_test), expected) diff --git a/tests/sparse/test_basic.py b/tests/sparse/test_basic.py index 18b7c08053..7749f83ca2 100644 --- a/tests/sparse/test_basic.py +++ b/tests/sparse/test_basic.py @@ -10,7 +10,7 @@ from pytensor import sparse from pytensor.compile.io import In from pytensor.configdefaults import config -from pytensor.gradient import GradientError +from pytensor.gradient import DisconnectedType, GradientError from pytensor.graph.basic import Apply from pytensor.graph.op import Op from pytensor.graph.replace import vectorize_graph @@ -950,6 +950,38 @@ def setup_method(self): self.op_class = EnsureSortedIndices self.op = ensure_sorted_indices + @pytest.mark.parametrize( + "format, constructor, shape", + [ + ("csr", scipy_sparse.csr_matrix, (2, 3)), + ("csc", scipy_sparse.csc_matrix, (3, 2)), + ], + ids=["csr", "csc"], + ) + def test_inplace(self, format, constructor, shape): + x = sparse.matrix(format=format, name="x", dtype="float64") + value = constructor( + ( + np.array([2.0, 3.0, 0.0]), + np.array([2, 0, 1], dtype="int32"), + np.array([0, 3, 3], dtype="int32"), + ), + shape=shape, + ) + expected = value.sorted_indices() + fn = pytensor.function( + [In(x, mutable=True)], + self.op_class(inplace=True)(x), + accept_inplace=True, + ) + result = fn(value) + # Also verify that the original input was mutated. + np.testing.assert_array_equal(value.data, expected.data) + np.testing.assert_array_equal(value.indices, expected.indices) + np.testing.assert_array_equal(result.data, expected.data) + np.testing.assert_array_equal(result.indices, expected.indices) + np.testing.assert_array_equal(result.indptr, expected.indptr) + def test_op(self): for format in sparse.sparse_formats: for shape in zip(range(5, 9), range(3, 7)[::-1], strict=True): @@ -976,6 +1008,105 @@ def test_grad(self): verify_grad_sparse(self.op, data, structured=False) +class TestHasCanonicalFormat(utt.InferShapeTester): + @pytest.mark.parametrize( + "format, constructor, shape", + [ + ("csr", scipy_sparse.csr_matrix, (2, 3)), + ("csc", scipy_sparse.csc_matrix, (3, 2)), + ], + ) + def test_op(self, format, constructor, shape): + x = sparse.matrix(format, name="x") + result = x.has_canonical_format + utt.assert_equal_computations([result], [sparse.has_canonical_format(x)]) + fn = pytensor.function([x], result) + + # The output is 0d boolean + assert result.type == pt.scalar(dtype="bool").type + + data = np.array([0, 2, 3], dtype=x.dtype) + indptr = np.array([0, 3, 3], dtype="int32") + indices_seq = [ + np.array(value, dtype="int32") + for value in ( + [0, 1, 2], # Unique and sorted + [2, 0, 1], # Unique but not sorted + [0, 0, 2], # Not unique + ) + ] + bools_seq = (True, False, False) + + for indices, expected in zip(indices_seq, bools_seq): + value = constructor((data, indices, indptr), shape=shape, copy=True) + assert fn(value) == expected + np.testing.assert_array_equal(value.indices, indices) + np.testing.assert_array_equal(value.data, data) + np.testing.assert_array_equal(value.indptr, indptr) + + # An empty matrix has canonical format + assert fn(constructor(shape, dtype=x.dtype)) + + def test_grad(self): + x = sparse.csr_matrix("x") + grad = pytensor.grad( + x.has_canonical_format, + x, + disconnected_inputs="ignore", + return_disconnected="disconnected", + ) + assert isinstance(grad.type, DisconnectedType) + + +class TestSumDuplicates(utt.InferShapeTester): + @pytest.mark.parametrize("dtype", ["float32", "float64"]) + @pytest.mark.parametrize( + "format, constructor, shape", + [ + ("csr", scipy_sparse.csr_matrix, (2, 3)), + ("csc", scipy_sparse.csc_matrix, (3, 2)), + ], + ) + def test_op(self, format, constructor, shape, dtype): + x = sparse.matrix(format=format, name="x", dtype=dtype) + value = constructor( + ( + np.array([2.0, 3.0, -2.0, 0.0], dtype=dtype), + np.array([2, 0, 2, 1], dtype="int32"), + np.array([0, 4, 4], dtype="int32"), + ), + shape=shape, + ) + original = value.copy() + fn = pytensor.function([x], sparse.sum_duplicates(x)) + result = fn(value) + + assert result.shape == shape + assert result.format == format + assert result.dtype == dtype + np.testing.assert_array_equal(result.data, [3.0, 0.0, 0.0]) + np.testing.assert_array_equal(result.indices, [0, 1, 2]) + np.testing.assert_array_equal(result.indptr, [0, 3, 3]) + np.testing.assert_array_equal(value.data, original.data) + np.testing.assert_array_equal(value.indices, original.indices) + np.testing.assert_array_equal(value.indptr, original.indptr) + + @pytest.mark.parametrize("format", ["csr", "csc"]) + def test_infer_shape(self, format): + variable, data = sparse_random_inputs(format, shape=(3, 5)) + self._compile_and_check( + variable, [sparse.sum_duplicates(*variable)], data, sparse.SumDuplicates + ) + + def test_grad(self): + x = sparse.csr_matrix("x") + gz = sparse.csr_matrix("gz") + summed = sparse.sum_duplicates(x) + utt.assert_equal_computations( + [pytensor.grad(cost=None, wrt=x, known_grads={summed: gz})], [gz] + ) + + class TestClean(utt.InferShapeTester): def setup_method(self): super().setup_method() diff --git a/tests/sparse/test_math.py b/tests/sparse/test_math.py index 4b43d1aefe..81edfc56c3 100644 --- a/tests/sparse/test_math.py +++ b/tests/sparse/test_math.py @@ -70,10 +70,10 @@ def eval_outputs(outputs): class TestAddMul: - def test_AddSS(self): + def test_AddSparseSparse(self): self._testSS(add) - def test_AddSD(self): + def test_AddSparseDense(self): self._testSD(add) def test_AddDS(self): @@ -1128,7 +1128,7 @@ def test_mul_s_v(self): utt.assert_allclose(spmat.toarray() * mat, out.toarray()) -class TestStructuredAddSV: +class TestStructuredAddSparseDenseVector: def test_structured_add_s_v_grad(self): sp_types = {"csc": scipy_sparse.csc_matrix, "csr": scipy_sparse.csr_matrix} for format in ("csr", "csc"): @@ -1156,13 +1156,13 @@ def test_structured_add_s_v(self): ) -class TestAddSSData(utt.InferShapeTester): +class TestAddSparseSparseData(utt.InferShapeTester): x = {} a = {} def setup_method(self): super().setup_method() - self.op_class = psm.AddSSData + self.op_class = psm.AddSparseSparseData for format in sparse_formats: variable = getattr(pytensor.sparse, format + "_matrix") @@ -1560,8 +1560,8 @@ def test_structured_dot_batched_sparse_raises(self): @pytest.mark.parametrize( "build", [ - pytest.param(lambda a, d: add(a, d), id="AddSD"), - pytest.param(lambda a, d: multiply(a, d), id="SparseDenseMultiply"), + pytest.param(lambda a, d: add(a, d), id="AddSparseDense"), + pytest.param(lambda a, d: multiply(a, d), id="MultiplySparseDense"), ], ) def test_sparse_dense_ops_batched_dense_raises(self, build): diff --git a/tests/sparse/test_rewriting.py b/tests/sparse/test_rewriting.py index 3d6d4dbeee..3bf21ad7f5 100644 --- a/tests/sparse/test_rewriting.py +++ b/tests/sparse/test_rewriting.py @@ -80,7 +80,7 @@ def test_local_mul_s_d(): f = pytensor.function(inputs, smath.mul_s_d(*inputs), mode="CVM") assert not any( - isinstance(node.op, smath.SparseDenseMultiply) + isinstance(node.op, smath.MultiplySparseDense) for node in f.maker.fgraph.toposort() ) @@ -102,7 +102,7 @@ def test_local_mul_s_v(): f = pytensor.function(inputs, smath.mul_s_v(*inputs), mode="CVM") assert not any( - isinstance(node.op, smath.SparseDenseVectorMultiply) + isinstance(node.op, smath.MultiplySparseDenseVector) for node in f.maker.fgraph.toposort() ) @@ -121,7 +121,7 @@ def test_local_structured_add_s_v(): f = pytensor.function(inputs, smath.structured_add_s_v(*inputs), mode="CVM") assert not any( - isinstance(node.op, smath.StructuredAddSV) + isinstance(node.op, smath.StructuredAddSparseDenseVector) for node in f.maker.fgraph.toposort() )