week-8 More ufunc overrides

Blake Griffith

This week was spent implementing the API for overriding ufuncs. After a lengthy discussion, I decided that it would be best to proceed with implementing things using a duck typing mechanism, similar to the one Python uses, instead of multi-methods. Nathaniel Smith and Pauli Virtanen had a lot of great suggestions on the interface which has been totally changed since last week. __ufunc_override__ has become __numpy_ufunc__ and is now a method which takes arguments like:

def __numpy_ufunc__(self, ufunc, method, i, inputs, kwargs):
    ...

Where:

ufunc is the ufunc object that was called.
method is a string indicating which Ufunc method was called (one of "__call__", "reduce", "reduceat", "accumulate", "outer", "inner").
i is the index of self in inputs.
inputs is a tuple of the input arguments to the ufunc
kwargs is a dictionary containing the optional input arguments of the ufunc. The out argument is always contained in kwargs, if given.

The implementation is best described by the NEP:

If one of the input arguments implements __numpy_ufunc__ it is executed instead of the Ufunc.
If more than one of the input arguments implements __numpy_ufunc__, they are tried in the following order: subclasses before superclasses, otherwise left to right. The first __numpy_ufunc__ method returning something else than NotImplemented determines the return value of the Ufunc. This is to accommodate Python's own MRO behavior.
If all __numpy_ufunc__ methods of the input arguments return NotImplemented, a TypeError is raised.
If a __numpy_ufunc__ method raises an error, the error is propagated immediately.

Demo¶

There have been several models working in a limited fashion so far, but here is an example. I added this to scipy/sparse/base.py

def __numpy_ufunc__(*args, **kwargs): self = args[0] ufunc = args[1] method = args[2] i = args[3] ufunc_args = args[4] if ufunc == np.multiply and method == "__call__": other_args = tuple(j for j in ufunc_args if j is not ufunc_args[i]) return self.multiply(*other_args, **kwds) else: return NotImplemented

So __numpy_ufunc__ is now a method on all sparse matrices.

In [21]:

import scipy.sparse as sp
import numpy as np

a = np.random.randint(5, size=(3,3))
b = np.random.randint(5, size=(3,3))

asp = sp.csr_matrix(a)
bsp = sp.csr_matrix(b)

In [22]:

np.multiply(a,b)

Out[22]:

array([[ 2,  0, 12],
       [ 0,  0, 12],
       [ 3,  0,  9]])

In [23]:

np.multiply(asp, bsp).todense()

Out[23]:

matrix([[ 2,  0, 12],
        [ 0,  0, 12],
        [ 3,  0,  9]], dtype=int64)

In [24]:

np.multiply(asp,b)

Out[24]:

matrix([[ 2,  0, 12],
        [ 0,  0, 12],
        [ 3,  0,  9]], dtype=int64)

In [25]:

np.multiply(a, bsp)

Out[25]:

matrix([[ 2,  0, 12],
        [ 0,  0, 12],
        [ 3,  0,  9]], dtype=int64)

Demo¶

Comments