Understanding the difference between numpy.log and numpy.sum for custom array containers

I am creating a custom container for a NumPy array by following the instructions on the SciPy website. I've written the following code which:

  • creates a class, NpContainer
  • defines custom behavior for the functions np.sum and np.log to output a string.
import numpy as np

HANDLED_FUNCTIONS = {}

class NpContainer:
    def __init__(self, val):
        self.val = val

    def __array__(self):
        return np.array(self.val)

    def __array_function__(self, func, types, args, kwargs):
        if func not in HANDLED_FUNCTIONS:
            raise NotImplementedError()

        return HANDLED_FUNCTIONS[func](*args, **kwargs)

def implements(np_function):
    def decorator(func):
        HANDLED_FUNCTIONS[np_function] = func
        return func
    return decorator

@implements(np.sum)
def sum(a, **kwargs):
    return 'Sum Val: {}'.format(np.sum(a.val, **kwargs))

@implements(np.log)
def log(a, **kwargs):
    return 'Log Val: {}'.format(np.log(a, **kwargs))

I test the code using:

if __name__ == "__main__":
    container1 = NpContainer(val=np.array([1., 2.]))

    sum_result = np.sum(container1)
    print(sum_result)
    print(type(sum_result))

    log_result = np.log(container1)
    print(log_result)
    print(type(log_result))

Sum produces the expected result.

Sum Val: 3.0
<class 'str'>

However np.log returns a NumPy array instead of a string.

[0.         0.69314718]
<class 'numpy.ndarray'>

Does anyone know why np.log skips my custom-defined function? Any help is appreciated!

1 answer

  • answered 2020-02-16 18:51 Grzegorz Skibinski

    Ok, so I think - I know what is happening. You implemented function __array_function__, which should fit np.sum but for np.log you should do __array_ufunc__, since it is an universal function (https://docs.scipy.org/doc/numpy/reference/arrays.classes.html#numpy.class.array_ufunc).

    There's some more flavour to it, which I recommend you to check out here:

    https://numpy.org/neps/nep-0018-array-function-protocol.html

    Now oddly enough- once __array_ufunc__ is implemented, even np.sum will be processed as ufunc, which messes up with the decorator.

    ufunc auto-casted function for np.sum is np.add, so the below will do the trick for you - although I would rather recommend to implement sum() as a function for NpContainer - so you can do container1.sum() instead

    import numpy as np
    
    HANDLED_FUNCTIONS = {}
    
    class NpContainer:
        def __init__(self, val):
            self.val = val
    
        def __array__(self):
            return np.array(self.val)
    
        def __array_function__(self, func, types, args, kwargs):
            if func not in HANDLED_FUNCTIONS:
                raise NotImplementedError()
    
            return HANDLED_FUNCTIONS[func](*args, **kwargs)
    
        def __array_ufunc__(self, ufunc, method, *args, **kwargs):
            if ufunc not in HANDLED_FUNCTIONS:
                raise NotImplementedError()
    
            return HANDLED_FUNCTIONS[ufunc](*args, **kwargs)        
    
    def implements(np_function):
        def decorator(func):
            HANDLED_FUNCTIONS[np_function] = func
            return func
        return decorator
    
    @implements(np.add)
    def sum(a, **kwargs):
        return 'Sum Val: {}'.format(np.sum(a.val, **kwargs))
    
    @implements(np.log)
    def log(a, **kwargs):
        return 'Log Val: {}'.format(np.log(a.val, **kwargs))
    
    if __name__ == "__main__":
        container1 = NpContainer(val=np.array([1., 2.]))
    
        log_result = np.log(container1)
        print(log_result)
        print(type(log_result))
    
        sum_result = np.sum(container1)
        print(sum_result)
        print(type(sum_result))