How to overide numpy ufunc with __array_ufunc__

Question

I'm trying to implement numpy's ufunc to work with a class, using the __array_ufunc__ method introduced in numpy v1.13.

To simplify, here's what the class could look like :

class toto():
    def __init__(self, value, name):
        self.value = value
        self.name = name
    def __add__(self, other):
    """add values and concatenate names"""
        return toto(self.value + other.value, self.name + other.name)
    def __sub__(self, other):
    """sub values and concatenate names"""
        return toto(self.value - other.value, self.name + other.name)  

tata = toto(5, "first")  
titi = toto(1, "second")

Now if I try to apply np.add between these two, I get the expected result, as np.add relies on add. But if I call say np.exp, I get an error as expected :

>>> np.exp(tata)
AttributeError: 'toto' object has no attribute 'exp'

Now what I would like to do is to "override" all numpy ufuncs to work smoothly with this class without having to redefine every methods (exp(self), log(self), ...) in the class.

I was planning to use numpy ufunc's [__array_ufunc__]1 to do this, but I don't really understand the doc as it doesn't provide a simple exemple of implementation.

If anyone has had any experience with this new functionnality that looks promising, could you provide a simple example ?

hpaulj · Answer 1 · 2019-03-28T02:05:55.227

If I extend your class with a __array_ufunc__ method (and __repr__):

class toto():
    def __init__(self, value, name):
        self.value = value
        self.name = name
    def __add__(self, other):
        """add values and concatenate names"""
        return toto(self.value + other.value, self.name + other.name)
    def __sub__(self, other):
        """sub values and concatenate names"""
        return toto(self.value - other.value, self.name + other.name)

    def __repr__(self):
        return f"toto: {self.value}, {self.name}"
    def __array_ufunc__(self, *args, **kwargs):
        print(args)
        print(kwargs)

And try some ufunc calls:

In [458]: np.exp(tata)                                                          
(<ufunc 'exp'>, '__call__', toto: 5, first)
{}
In [459]: np.exp.reduce(tata)                                                   
(<ufunc 'exp'>, 'reduce', toto: 5, first)
{}
In [460]: np.multiply.reduce(tata)                                              
(<ufunc 'multiply'>, 'reduce', toto: 5, first)
{}
In [461]: np.exp.reduce(tata,axes=(1,2))                                        
(<ufunc 'exp'>, 'reduce', toto: 5, first)
{'axes': (1, 2)}
In [463]: np.exp.reduce(tata,axes=(1,2),out=np.arange(3))                       
(<ufunc 'exp'>, 'reduce', toto: 5, first)
{'axes': (1, 2), 'out': (array([0, 1, 2]),)}

That shows the information that your class receives. Evidently you can do what you want that. It can return NotImplemented. I suppose in your case it could apply the first argument to your self.value, or do some custom calculation.

For example if I add

      val = args[0].__call__(self.value) 
      return toto(val, self.name)

I get:

In [468]: np.exp(tata)                                                          
(<ufunc 'exp'>, '__call__', toto: 5, first)
{}
Out[468]: toto: 148.4131591025766, first
In [469]: np.sin(tata)                                                          
(<ufunc 'sin'>, '__call__', toto: 5, first)
{}
Out[469]: toto: -0.9589242746631385, first

However if I put the object in an array, I still get the method error

In [492]: np.exp(np.array(tata))                                                
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-492-4dc37eb0ffe3> in <module>
----> 1 np.exp(np.array(tata))

AttributeError: 'toto' object has no attribute 'exp'

Apparently ufunc on an object dtype array iterates on the elements of the array, expecting to use a 'relevant' method. For np.add (+) it looks for the __add__ method. For np.exp it looks for an exp method. This __array_ufunc__ isn't called.

So it looks like it's intended more for a subclass of ndarray, or something equivalent. You, I think, are trying to implement a class that can work as elements of an object dtype array.

Thanks for your insights. Although they do mention in the [doc](https://docs.scipy.org/doc/numpy/reference/arrays.classes.html#numpy.class.__array_ufunc__) that this method can be defined to extend any class, ndarray sub-classed or not. — mocquin, Mar 28 '19 at 10:59

score 0 · Answer 2 · answered Mar 24 '23 at 15:49

I think you are missing the __array_function__ protocol, https://numpy.org/neps/nep-0018-array-function-protocol.html

__array_ufunc__ will work only with certain numpy ufuncs but not with all. When not available, numpy will dispatch with the __array_function__ protocol, https://numpy.org/devdocs/release/1.17.0-notes.html#numpy-functions-now-always-support-overrides-with-array-function

A simple example follows:

import numpy as np
import logging
import inspect

HANDLED_FUNCTIONS = {}

def implements(numpy_function):
    """Register an __array_function__ implementation for MyArray objects."""
    def decorator(func):
        HANDLED_FUNCTIONS[numpy_function] = func
        return func
    return decorator

class MyArray(object):
    def __array_function__(self, func, types, args, kwargs):
        logging.debug('{} {}'.format(inspect.currentframe().f_code.co_name, func))
        if func not in HANDLED_FUNCTIONS:
            return NotImplemented
        if not all(issubclass(t, MyArray) for t in types):
            return NotImplemented
        return HANDLED_FUNCTIONS[func](*args, **kwargs)
    
    def __array_ufunc__(self, ufunc, method, inputs, *args, **kwargs):
        logging.debug('{} {}'.format(inspect.currentframe().f_code.co_name, ufunc))
        if ufunc not in HANDLED_FUNCTIONS:
            return NotImplemented
        out = kwargs.pop('out', None)
    
        if out is not None:
            HANDLED_FUNCTIONS[ufunc](inputs, *args, out=out[0], **kwargs)
            return
        else:
            return HANDLED_FUNCTIONS[ufunc](inputs, *args, out=None, **kwargs)
        
    def __init__(self, inlist):
        self.list = inlist[:]        
    
    @property
    def ndim(self):
        return 1

    @property
    def shape(self):
        return (len(self.list), )
    
    @property
    def dtype(self):
        return np.dtype(np.int32)

    def __str__(self):
        return "MyArray " + str(self.list)
    
    def __add__(self, other, *args, **kwargs):
        logging.debug('{}'.format(inspect.currentframe().f_code.co_name))
        return self.add(other, *args, **kwargs)

    @implements(np.add)
    def add(self, *args, **kwargs):
        strng = "MyClass add, out {} {}".format( kwargs.get('out', None), len(args) )
        logging.debug('{} {}'.format(inspect.currentframe().f_code.co_name, strng))
        out = kwargs.get('out', None)
        if out is None:
            return MyArray([el + args[0] for el in self.list])
        else:
            for i,el in enumerate(self.list):
                out[i] = args[0] + el

    # implements np.sum is required when one wants to use the np.sum on this object            
    @implements(np.sum)
    def sum(self, *args, **kwargs):
        return sum(self.list)  # return self.list.ndim    

def main():
    logging.basicConfig(level=logging.DEBUG)

    A = MyArray(np.array([1,2]))
    
    # test np.sum
    print ("sum" , np.sum(A, axis=1))
    
    # test add
    B = A.add(2)
    printit(B, 'B')
    
    out = MyArray([20,30])
    printit(out,'out')
    A.add(2,out=out)
    printit(out,'out')

    # test np.add
    # see comments on __add__ 
    #B = A+2
    B = np.add(A,2)
    printit(B, 'B')

    B = A+2
    printit(B, 'B')

    np.add(A,2,out=out)
    printit(out, "out")

How to overide numpy ufunc with __array_ufunc__

2 Answers2