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):


  • 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.


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]:
array([[ 2,  0, 12],
       [ 0,  0, 12],
       [ 3,  0,  9]])
In [23]:
np.multiply(asp, bsp).todense()
matrix([[ 2,  0, 12],
        [ 0,  0, 12],
        [ 3,  0,  9]], dtype=int64)
In [24]:
matrix([[ 2,  0, 12],
        [ 0,  0, 12],
        [ 3,  0,  9]], dtype=int64)
In [25]:
np.multiply(a, bsp)
matrix([[ 2,  0, 12],
        [ 0,  0, 12],
        [ 3,  0,  9]], dtype=int64)


Comments powered by Disqus