Skip to content

Commit

Permalink
Merge pull request #29 from projekter/empty-product
Browse files Browse the repository at this point in the history
Fix #28: Empty product
  • Loading branch information
asistradition authored Nov 7, 2024
2 parents 5ea1658 + 4f168eb commit 2ef0643
Show file tree
Hide file tree
Showing 8 changed files with 133 additions and 79 deletions.
11 changes: 8 additions & 3 deletions sparse_dot_mkl/_dense_dense.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@ def _dense_matmul(matrix_a, matrix_b, scalar=1., out=None, out_scalar=None):
# The complex versions of these functions take void pointers instead of passed structs
# So create a C struct if necessary to be passed by reference
scalar = _mkl_scalar(scalar, complex_type, double_precision)
out_scalar = _mkl_scalar(out_scalar, complex_type, double_precision)
out_scalar = _mkl_scalar(0 if out is None else out_scalar, complex_type, double_precision)

func(layout_a,
111,
Expand All @@ -75,8 +75,13 @@ def _dense_dot_dense(matrix_a, matrix_b, cast=False, scalar=1., out=None, out_sc
# Check for edge condition inputs which result in empty outputs
if _empty_output_check(matrix_a, matrix_b):
debug_print("Skipping multiplication because A (dot) B must yield an empty matrix")
final_dtype = np.float64 if matrix_a.dtype != matrix_b.dtype or matrix_a.dtype != np.float32 else np.float32
return _out_matrix((matrix_a.shape[0], matrix_b.shape[1]), final_dtype, out_arr=out)
output_arr = _out_matrix((matrix_a.shape[0], matrix_b.shape[1]),
_type_check(matrix_a, matrix_b, cast=cast, convert=False), out_arr=out)
if out is None or (out_scalar is not None and not out_scalar):
output_arr.fill(0)
elif out_scalar is not None:
output_arr *= out_scalar
return output_arr

matrix_a, matrix_b = _type_check(matrix_a, matrix_b, cast=cast)

Expand Down
78 changes: 50 additions & 28 deletions sparse_dot_mkl/_gram_matrix.py
Original file line number Diff line number Diff line change
Expand Up @@ -93,7 +93,7 @@ def _gram_matrix_sparse(


# Dict keyed by ('double_precision_bool', 'complex_bool')
_mkl_skryd_funcs = {
_mkl_syrkd_funcs = {
(False, False): MKL._mkl_sparse_s_syrkd,
(True, False): MKL._mkl_sparse_d_syrkd,
(False, True): MKL._mkl_sparse_c_syrkd,
Expand Down Expand Up @@ -129,19 +129,29 @@ def _gram_matrix_sparse_to_dense(
_order_mkl_handle(sp_ref_a)

out_dtype = _output_dtypes[(double_prec, complex_type)]
func = _mkl_skryd_funcs[(double_prec, complex_type)]
func = _mkl_syrkd_funcs[(double_prec, complex_type)]

out_dim = matrix_a.shape[0 if aat else 1]

output_arr = _out_matrix(
(out_dim, out_dim),
out_dtype,
order="C", out_arr=out)
order="C",
out_arr=out,
initialize_zeros=True
)
_, output_ld = _get_numpy_layout(output_arr)

if _empty_output_check(matrix_a, matrix_a):
_destroy_mkl_handle(sp_ref_a)
if out is None or (out_scalar is not None and not out_scalar):
output_arr.fill(0)
elif out_scalar is not None:
output_arr *= out_scalar
return output_arr

if out is None:
out_scalar = 0.

scalar = _mkl_scalar(scalar, complex_type, double_prec)
out_scalar = _mkl_scalar(out_scalar, complex_type, double_prec)
Expand All @@ -165,7 +175,7 @@ def _gram_matrix_sparse_to_dense(
# matrix. This stupid thing only happens with specific flags
# I could probably leave it but it's pretty annoying

if not aat and out is None and not complex_type:
if not aat and out is None:
output_arr[np.tril_indices(output_arr.shape[0], k=-1)] = 0.0

return output_arr
Expand Down Expand Up @@ -208,6 +218,13 @@ def _gram_matrix_dense_to_dense(
:rtype: numpy.ndarray
"""

if aat and np.iscomplexobj(matrix_a):
raise ValueError(
"transpose=True with dense complex data currently "
"fails with an Intel oneMKL ERROR: "
"Parameter 3 was incorrect on entry to cblas_csyrk"
)

# Get dimensions
n, k = matrix_a.shape if aat else matrix_a.shape[::-1]

Expand All @@ -223,15 +240,16 @@ def _gram_matrix_dense_to_dense(
(n, n),
out_dtype,
order="C" if layout_a == LAYOUT_CODE_C else "F",
out_arr=out
out_arr=out,
initialize_zeros=True
)

# The complex versions of these functions take void pointers instead of
# passed structs, so create a C struct if necessary to be passed by
# reference
scalar = _mkl_scalar(scalar, complex_type, double_precision)
out_scalar = _mkl_scalar(out_scalar, complex_type, double_precision)

func(
layout_a,
MKL_UPPER,
Expand All @@ -243,7 +261,7 @@ def _gram_matrix_dense_to_dense(
ld_a,
out_scalar if not complex_type else _ctypes.byref(scalar),
output_arr,
n,
n
)

return output_arr
Expand Down Expand Up @@ -279,6 +297,22 @@ def _gram_matrix(
:rtype: scipy.sparse.csr_matrix, np.ndarray
"""

if _sps.issparse(matrix) and not (is_csr(matrix) or is_csc(matrix)):
raise ValueError(
"gram_matrix requires sparse matrix to be CSR or CSC format"
)
elif is_csc(matrix) and not cast:
raise ValueError(
"gram_matrix cannot use a CSC matrix unless cast=True"
)
elif out is not None and not dense:
raise ValueError(
"out argument cannot be used with sparse (dot) sparse "
"matrix multiplication"
)
elif out is not None and not isinstance(out, np.ndarray):
raise ValueError("out argument must be dense")

# Check for edge condition inputs which result in empty outputs
if _empty_output_check(matrix, matrix):
debug_print(
Expand All @@ -290,25 +324,18 @@ def _gram_matrix(
if transpose
else (matrix.shape[0], matrix.shape[0])
)
output_func = _sps.csr_matrix if _sps.isspmatrix(matrix) else np.zeros
return output_func(output_shape, dtype=matrix.dtype)

if np.iscomplexobj(matrix):
raise ValueError(
"gram_matrix_mkl does not support complex datatypes"
)
if out is None:
output_func = np.zeros if dense else _sps.csr_matrix
return output_func(output_shape, dtype=matrix.dtype)
elif out_scalar is not None and not out_scalar:
out.fill(0)
elif out_scalar is not None:
out *= out_scalar
return out

matrix = _type_check(matrix, cast=cast)

if _sps.issparse(matrix) and not (is_csr(matrix) or is_csc(matrix)):
raise ValueError(
"gram_matrix requires sparse matrix to be CSR or CSC format"
)
elif is_csc(matrix) and not cast:
raise ValueError(
"gram_matrix cannot use a CSC matrix unless cast=True"
)
elif not _sps.issparse(matrix):
if not _sps.issparse(matrix):
return _gram_matrix_dense_to_dense(
matrix,
aat=transpose,
Expand All @@ -322,11 +349,6 @@ def _gram_matrix(
out=out,
out_scalar=out_scalar
)
elif out is not None:
raise ValueError(
"out argument cannot be used with sparse (dot) sparse "
"matrix multiplication"
)
else:
return _gram_matrix_sparse(
matrix,
Expand Down
37 changes: 24 additions & 13 deletions sparse_dot_mkl/_mkl_interface/_common.py
Original file line number Diff line number Diff line change
Expand Up @@ -770,12 +770,14 @@ def _is_valid_dtype(matrix, complex_dtype=False, all_dtype=False):
return matrix.dtype in NUMPY_FLOAT_DTYPES


def _type_check(matrix_a, matrix_b=None, cast=False, allow_complex=True):
def _type_check(matrix_a, matrix_b=None, cast=False, allow_complex=True, convert=True):
"""
Make sure that both matrices are single precision floats or both are
double precision floats.
If not, convert to double precision floats if cast is True,
or raise an error if cast is False
If convert is set to False, the resulting data type is returned without
any conversion happening.
"""

_n_complex = _np.iscomplexobj(matrix_a) + _np.iscomplexobj(matrix_b)
Expand All @@ -785,17 +787,17 @@ def _type_check(matrix_a, matrix_b=None, cast=False, allow_complex=True):

# If there's no matrix B and matrix A is valid dtype, return it
if matrix_b is None and _is_valid_dtype(matrix_a, all_dtype=True):
return matrix_a
return matrix_a if convert else matrix_a.dtype

# If matrix A is complex but not csingle or cdouble, and cast is True,
# convert it to a cdouble
elif matrix_b is None and cast and _n_complex == 1:
return _cast_to(matrix_a, _np.cdouble)
return _cast_to(matrix_a, _np.cdouble) if convert else _np.cdouble

# If matrix A is real but not float32 or float64, and cast is True,
# convert it to a float64
elif matrix_b is None and cast:
return _cast_to(matrix_a, _np.float64)
return _cast_to(matrix_a, _np.float64) if convert else _np.float64

# Raise an error - the dtype is invalid and cast is False
elif matrix_b is None:
Expand All @@ -809,7 +811,7 @@ def _type_check(matrix_a, matrix_b=None, cast=False, allow_complex=True):
_is_valid_dtype(matrix_a, all_dtype=True) and
matrix_a.dtype == matrix_b.dtype
):
return matrix_a, matrix_b
return (matrix_a, matrix_b) if convert else matrix_a.dtype

# If neither matrix is complex and cast is True, convert to float64s
# and return them
Expand All @@ -818,7 +820,7 @@ def _type_check(matrix_a, matrix_b=None, cast=False, allow_complex=True):
f"Recasting matrix data types {matrix_a.dtype} and "
f"{matrix_b.dtype} to np.float64"
)
return _cast_to(matrix_a, _np.float64), _cast_to(matrix_b, _np.float64)
return (_cast_to(matrix_a, _np.float64), _cast_to(matrix_b, _np.float64)) if convert else _np.float64

# If both matrices are complex and cast is True, convert to cdoubles
# and return them
Expand All @@ -827,7 +829,7 @@ def _type_check(matrix_a, matrix_b=None, cast=False, allow_complex=True):
f"Recasting matrix data types {matrix_a.dtype} and "
f"{matrix_b.dtype} to _np.cdouble"
)
return _cast_to(matrix_a, _np.cdouble), _cast_to(matrix_b, _np.cdouble)
return (_cast_to(matrix_a, _np.cdouble), _cast_to(matrix_b, _np.cdouble)) if convert else _np.cdouble

# Cast reals and complex matrices together
elif (
Expand All @@ -838,7 +840,7 @@ def _type_check(matrix_a, matrix_b=None, cast=False, allow_complex=True):
debug_print(
f"Recasting matrix data type {matrix_b.dtype} to {matrix_a.dtype}"
)
return matrix_a, _cast_to(matrix_b, matrix_a.dtype)
return (matrix_a, _cast_to(matrix_b, matrix_a.dtype)) if convert else matrix_a.dtype

elif (
cast and
Expand All @@ -848,14 +850,14 @@ def _type_check(matrix_a, matrix_b=None, cast=False, allow_complex=True):
debug_print(
f"Recasting matrix data type {matrix_a.dtype} to {matrix_b.dtype}"
)
return _cast_to(matrix_a, matrix_b.dtype), matrix_b
return (_cast_to(matrix_a, matrix_b.dtype), matrix_b) if convert else matrix_b.dtype

elif cast and _n_complex == 1:
debug_print(
f"Recasting matrix data type {matrix_a.dtype} and {matrix_b.dtype}"
f" to np.cdouble"
)
return _cast_to(matrix_a, _np.cdouble), _cast_to(matrix_b, _np.cdouble)
return (_cast_to(matrix_a, _np.cdouble), _cast_to(matrix_b, _np.cdouble)) if convert else _np.cdouble

# If cast is False, can't cast anything together
elif not cast:
Expand All @@ -882,9 +884,16 @@ def _mkl_scalar(scalar, complex_type, double_precision):
return float(scalar)


def _out_matrix(shape, dtype, order="C", out_arr=None, out_t=False):
def _out_matrix(
shape,
dtype,
order="C",
out_arr=None,
out_t=False,
initialize_zeros=False
):
"""
Create an all-zero matrix or check to make sure that
Create an undefined matrix or check to make sure that
the provided output array matches
:param shape: Required output shape
Expand All @@ -904,8 +913,10 @@ def _out_matrix(shape, dtype, order="C", out_arr=None, out_t=False):
out_t = False if out_t is None else out_t

# If there's no output array allocate a new array and return it
if out_arr is None:
if out_arr is None and initialize_zeros:
return _np.zeros(shape, dtype=dtype, order=order)
elif out_arr is None:
return _np.ndarray(shape, dtype=dtype, order=order)

# Check and make sure the order is correct
# Note 1d arrays have both flags set
Expand Down
17 changes: 9 additions & 8 deletions sparse_dot_mkl/_sparse_dense.py
Original file line number Diff line number Diff line change
Expand Up @@ -106,7 +106,7 @@ def _sparse_dense_matmul(

# Create a C struct if necessary to be passed
scalar = _mkl_scalar(scalar, cplx, dbl)
out_scalar = _mkl_scalar(out_scalar, cplx, dbl)
out_scalar = _mkl_scalar(0 if out is None else out_scalar, cplx, dbl)

ret_val = func(
11 if transpose else 10,
Expand Down Expand Up @@ -165,14 +165,15 @@ def _sparse_dot_dense(
debug_print(
"Skipping multiplication because A (dot) B must yield empty matrix"
)
final_dtype = (
np.float64
if matrix_a.dtype != matrix_b.dtype or matrix_a.dtype != np.float32
else np.float32
)
return _out_matrix(
(matrix_a.shape[0], matrix_b.shape[1]), final_dtype, out_arr=out
output_arr = _out_matrix(
(matrix_a.shape[0], matrix_b.shape[1]),
_type_check(matrix_a, matrix_b, cast=cast, convert=False), out_arr=out
)
if out is None or (out_scalar is not None and not out_scalar):
output_arr.fill(0)
elif out_scalar is not None:
output_arr *= out_scalar
return output_arr

matrix_a, matrix_b = _type_check(matrix_a, matrix_b, cast=cast)

Expand Down
9 changes: 6 additions & 3 deletions sparse_dot_mkl/_sparse_sparse.py
Original file line number Diff line number Diff line change
Expand Up @@ -169,16 +169,19 @@ def _sparse_dot_sparse(
# Check for edge condition inputs which result in empty outputs
if _empty_output_check(matrix_a, matrix_b):

final_dtype = _type_check(matrix_a, matrix_b, cast=cast, convert=False)
if dense:
return _out_matrix(
output_arr = _out_matrix(
(matrix_a.shape[0], matrix_b.shape[1]),
matrix_a.dtype,
final_dtype,
out_arr=out
)
output_arr.fill(0)
return output_arr
else:
return default_output(
(matrix_a.shape[0], matrix_b.shape[1]),
dtype=matrix_a.dtype
dtype=final_dtype
)

# Check dtypes
Expand Down
4 changes: 2 additions & 2 deletions sparse_dot_mkl/_sparse_sypr.py
Original file line number Diff line number Diff line change
Expand Up @@ -69,8 +69,8 @@ def _sypr_sparse_A_dense_B(
matrix_b,
layout_b,
ld_b,
float(out_scalar) if a_scalar is not None else 1.0,
float(out_scalar) if out_scalar is not None else 1.0,
float(a_scalar) if a_scalar is not None else 1.0,
0.0 if out is None else (float(out_scalar) if out_scalar is not None else 1.0),
output_arr,
output_layout,
output_ld,
Expand Down
Loading

0 comments on commit 2ef0643

Please sign in to comment.