-
Notifications
You must be signed in to change notification settings - Fork 2
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add explicit backends.py files in dask.array/dataframe
This collects lazily registered functions for cupy, sparse, scipy, and cudf.
- Loading branch information
Showing
8 changed files
with
73 additions
and
68 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,33 @@ | ||
from .core import tensordot_lookup, concatenate_lookup | ||
|
||
|
||
@tensordot_lookup.register_lazy('cupy') | ||
@concatenate_lookup.register_lazy('cupy') | ||
def register_cupy(): | ||
import cupy | ||
concatenate_lookup.register(cupy.ndarray, cupy.concatenate) | ||
tensordot_lookup.register(cupy.ndarray, cupy.tensordot) | ||
|
||
|
||
@tensordot_lookup.register_lazy('sparse') | ||
@concatenate_lookup.register_lazy('sparse') | ||
def register_sparse(): | ||
import sparse | ||
concatenate_lookup.register(sparse.COO, sparse.concatenate) | ||
tensordot_lookup.register(sparse.COO, sparse.tensordot) | ||
|
||
|
||
@concatenate_lookup.register_lazy('scipy') | ||
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) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,37 @@ | ||
from .methods import concat_dispatch | ||
from .core import get_parallel_type, meta_nonempty, make_meta | ||
|
||
|
||
###################################### | ||
# cuDF: Pandas Dataframes on the GPU # | ||
###################################### | ||
|
||
|
||
@concat_dispatch.register_lazy('cudf') | ||
@get_parallel_type.register_lazy('cudf') | ||
@meta_nonempty.register_lazy('cudf') | ||
@make_meta.register_lazy('cudf') | ||
def _register_cudf(): | ||
import cudf | ||
import dask_cudf | ||
get_parallel_type.register(cudf.DataFrame, lambda _: dask_cudf.DataFrame) | ||
get_parallel_type.register(cudf.Series, lambda _: dask_cudf.Series) | ||
get_parallel_type.register(cudf.Index, lambda _: dask_cudf.Index) | ||
|
||
@meta_nonempty.register((cudf.DataFrame, cudf.Series, cudf.Index)) | ||
def _(x): | ||
y = meta_nonempty(x.to_pandas()) # TODO: add iloc[:5] | ||
return cudf.from_pandas(y) | ||
|
||
@make_meta.register((cudf.Series, cudf.DataFrame)) | ||
def _(x): | ||
return x.head(0) | ||
|
||
@make_meta.register(cudf.Index) | ||
def _(x): | ||
return x[:0] | ||
|
||
concat_dispatch.register( | ||
(cudf.DataFrame, cudf.Series, cudf.Index), | ||
cudf.concat | ||
) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters