Skip to content

A catch-all protocol for numpy-like duck arrays #11129

@mrocklin

Description

@mrocklin

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.

cc @shoyer @hameerabbasi @njsmith @ericmjl

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions