Numpy use dispatcher for argument displacement

Viewed 262

This question is specific to NEP 18, which proposes a dispatch mechanism for using an arbitrary functions together with classes that implement the __array_function__ interface (such as np.ndarray).

def _einsum_dispatcher(subscripts, *operands):
    return (*operands,), dict(subscripts=subscripts)

@np.core.overrides.array_function_dispatch(_einsum_dispatcher, verify=False)
def my_einsum(*operands, **subscripts):
    raise RuntimeError("This should currently never happen")

class ArrayLike:
    def __array_function__(self, func, types, args, kwargs):
        return f"""
               func = {func},
               types = {types},
               args = {args},
               kwargs = {kwargs}
               """

Running a test for my_einsum we get

a = ArrayLike()
print(my_einsum('...', a))

Output:

RuntimeError: This should currently never happen

The function never gets intercepted by __array_function__ method for ArrayLike. Although, this can be solved by changing the dispatcher a bit:

def _einsum_dispatcher(subscripts, *operands):
    return (*operands,)

Now when testing my_einsum the exact same way again:

a = ArrayLike()
print(my_einsum('...', a))

Output:

func = <function my_einsum at 0x7efded61bd30>,
types = (<class '__main__.ArrayLike'>,),
args = ('...', <__main__.ArrayLike object at 0x7efdef83efa0>),
kwargs = {}

The function is correct, but how do I move the subscripts argument into kwargs ? Can the array_function_dispatch only be used to dispatch argument types?

0 Answers
Related