week-8 More ufunc overrides

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.

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)

np.multiply(a,b)

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

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

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

np.multiply(asp,b)

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

np.multiply(a, bsp)

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