def register_scipy_sparse():
import scipy.sparse
def _concatenate(L, axis=0):
if axis == 0:
return scipy.sparse.vstack(L)
elif axis == 1:
return scipy.sparse.hstack(L)
else:
msg = ("Can only concatenate scipy sparse matrices for axis in "
"{0, 1}. Got %s" % axis)
raise ValueError(msg)
concatenate_lookup.register(scipy.sparse.spmatrix, _concatenate)
tensordot_lookup.register(scipy.sparse.spmatrix, _tensordot_scipy_sparse)
def register_cupy():
import cupy
from dask.array.dispatch import percentile_lookup
concatenate_lookup.register(cupy.ndarray, cupy.concatenate)
tensordot_lookup.register(cupy.ndarray, cupy.tensordot)
percentile_lookup.register(cupy.ndarray, percentile)
@einsum_lookup.register(cupy.ndarray)
def _cupy_einsum(*args, **kwargs):
# NB: cupy does not accept `order` or `casting` kwargs - ignore
kwargs.pop("casting", None)
kwargs.pop("order", None)
return cupy.einsum(*args, **kwargs)
def register_cupyx():
from cupyx.scipy.sparse import spmatrix
try:
from cupyx.scipy.sparse import hstack, vstack
except ImportError as e:
raise ImportError(
"Stacking of sparse arrays requires at least CuPy version 8.0.0"
) from e
def _concat_cupy_sparse(L, axis=0):
if axis == 0:
return vstack(L)
elif axis == 1:
return hstack(L)
else:
msg = ("Can only concatenate cupy sparse matrices for axis in "
"{0, 1}. Got %s" % axis)
raise ValueError(msg)
concatenate_lookup.register(spmatrix, _concat_cupy_sparse)
tensordot_lookup.register(spmatrix, _tensordot_scipy_sparse)
def register_sparse():
import sparse
concatenate_lookup.register(sparse.COO, sparse.concatenate)
tensordot_lookup.register(sparse.COO, sparse.tensordot)
import numpy as np
from dask.array.dispatch import (
concatenate_lookup,
divide_lookup,
einsum_lookup,
empty_lookup,
percentile_lookup,
tensordot_lookup,
)
from dask.array.numpy_compat import divide as np_divide
from dask.array.numpy_compat import ma_divide
from dask.array.percentile import _percentile
concatenate_lookup.register((object, np.ndarray), np.concatenate)
tensordot_lookup.register((object, np.ndarray), np.tensordot)
einsum_lookup.register((object, np.ndarray), np.einsum)
empty_lookup.register((object, np.ndarray), np.empty)
empty_lookup.register(np.ma.masked_array, np.ma.empty)
divide_lookup.register((object, np.ndarray), np_divide)
divide_lookup.register(np.ma.masked_array, ma_divide)
@percentile_lookup.register(np.ndarray)
def percentile(a, q, method="linear"):
return _percentile(a, q, method)
@concatenate_lookup.register(np.ma.masked_array)
def _concatenate(arrays, axis=0):
out = np.ma.concatenate(arrays, axis=axis)