-
-
Notifications
You must be signed in to change notification settings - Fork 11.1k
Description
There are several functions for which I would like to see protocols constructed. I've raised issues for #11074 and #11128 but these are just special cases of a much larger issue that includes many operations. The sense I've gotten is that the process to change numpy takes a while, so I'm inclined to find a catch-all solution that can serve as a catch-all while things evolve.
To that end I propose that duck-arrays include a method that returns a module that mimics the numpy namespace
class ndarray:
def __array_module__(self):
import numpy as np
return np
class DaskArray:
def __array_module__(self):
import dask.array as da
return da
class CuPyArray:
def __array_module__(self):
import cupy as cp
return cp
class SparseArray:
def __array_module__(self):
import sparse
return sparse
...
Then, in various functions like stack or concatenate we check for these modules
def stack(args, **kwargs):
modules = {arg.__array_module__() for arg in args}
if len(modules) == 1:
module = list(modules)[0]
if module != numpy:
return module.stack(args, **kwargs)
...
There are likely several things wrong the implementation above, but my hope is that it gets a general point across that we'll dispatch wholesale to the module of the provided duck arrays.