Trapz dx doesn't trigger array_function mechanism

[mocquin]

It seems that when using numpy's trapz numpy.trapzwith the __array_function__ mechanism, only y and x are caught, but not dx, with numpy's signature numpy.trapz(y, x=None, dx=1.0, axis=-1).

Reproducing code example:

Here is a sample code that creates a numerical-labelled object, basically a value (scalar or array) and a label as a string.
A wrapped version of numpy trapz is created and registered, so that anytime numpy recieves a NumericalLabelled in a trapz call, it relies on the wrapped version. But it seems that the wrapped version is never called when only dx is a NumericalLabelled. I would expect it is called as soon as any of the input is of the type NumericalLabelled.

import numpy as np

HANDLED_FUNCTIONS = {}

class NumericalLabeled():
    def __init__(self, value, label=""):
        self.value = value
        self.label = label
        
    def __repr__(self):
        return "NumericalLabelled<"+str(self.value) + "," + self.label+">"
    
    def __array_function__(self, func, types, args, kwargs):
        print("Got into array function")
        if func not in HANDLED_FUNCTIONS:
            return NotImplemented
        return HANDLED_FUNCTIONS[func](*args, **kwargs)
    
def make_numericallabelled(x, label=""):
    """
    Helper function to cast anything into a NumericalLabelled object.
    """
    if isinstance(x, NumericalLabeled):
        return x
    else:
        return NumericalLabeled(x, label=label)
    
# Numpy functions            
# Override functions - used with __array_function__
def implements(np_function):
    def decorator(func):
        HANDLED_FUNCTIONS[np_function] = func
        return func
    return decorator    
    
@implements(np.trapz)
def np_trapz(q, x=None, dx=1, **kwargs):
    """
    Numpy's trapz wrapper for NumericalLabelled.
    """
    # first convert q into a NumericalLabelled to use `q.value` 
    q = make_numericallabelled(q)
    if x is None:    
        # using dx.value and dx.label
        dx = make_numericallabelled(dx, label="dx")
        return NumericalLabeled(np.trapz(q.value, dx=dx.value, x=None, **kwargs),
                                q.label + dx.label,
                    )
    else:
        # using x/value and x.label
        x = make_numericallabelled(x, label="x")
        return NumericalLabeled(np.trapz(q.value, x=x.value, **kwargs),
                                q.label + x.label,
                    )

def main():
    # create a scalar to use as dx
    half = NumericalLabeled(0.5, "half")
    # create an array to use as x
    x = NumericalLabeled(np.arange(5), "x")
    # then 
    # this works
    print(np.trapz(NumericalLabeled(np.arange(5), "a")))
    # this also works
    print(np.trapz(np.arange(5), x=x))
    # but not this
    print(np.trapz(np.arange(5), dx=half))
    # TypeError: unsupported operand type(s) for *: 'NumericalLabeled' and 'int'
main()

Error message:

Got into array function
NumericalLabelled<8.0,adx>
Got into array function
NumericalLabelled<8.0,x>

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-30-3401bee8c135> in <module>
     68     np.trapz(np.arange(5), dx=half)
     69     # TypeError: unsupported operand type(s) for *: 'NumericalLabeled' and 'int'
---> 70 main()

<ipython-input-30-3401bee8c135> in main()
     66     np.trapz(np.arange(5), x=x)
     67     # but not this
---> 68     np.trapz(np.arange(5), dx=half)
     69     # TypeError: unsupported operand type(s) for *: 'NumericalLabeled' and 'int'
     70 main()

<__array_function__ internals> in trapz(*args, **kwargs)

/opt/anaconda3/lib/python3.8/site-packages/numpy/lib/function_base.py in trapz(y, x, dx, axis)
   4161     slice2[axis] = slice(None, -1)
   4162     try:
-> 4163         ret = (d * (y[tuple(slice1)] + y[tuple(slice2)]) / 2.0).sum(axis)
   4164     except ValueError:
   4165         # Operations didn't work, cast to ndarray

TypeError: unsupported operand type(s) for *: 'NumericalLabeled' and 'int'

NumPy/Python version information:

1.20.0 3.8.5 (default, Sep 4 2020, 02:22:02) [Clang 10.0.0 ]

[mhvk]

In general, the idea is to trigger only on arrays, and here the docs state dx is a scaler. However, the implementation definitely allows an array, so agreed this is a bug. Would be good to adjust the docstring at the same time.

[mocquin]

What is the underlying reason to trigger the array_function only on arrays ? There are lots of functions in numpy that can be used with only scalars, but hence can't be overloaded

[mhvk]

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

It is a conscious choice not to override arguments like axis or keepdims, as it is assumed the intent is that the functions continues to work the same, but just on a different array type. But probably most fruitful to discuss specific other examples where an override of a given argument would be good - maybe like your example here it is a bug that they are not overridden.

[mocquin]

Agreed on arguments like axis and keepdims. Another example is the random functions like numpy.random.normal but I am not sure it fails for the same reason :

import numpy as np

HANDLED_FUNCTIONS = {}

class NumericalLabeled():
    def __init__(self, value, label=""):
        self.value = value
        self.label = label
        
    def __repr__(self):
        return "NumericalLabelled<"+str(self.value) + "," + self.label+">"
    
    def __array_function__(self, func, types, args, kwargs):
        print("Got into array function")
        if func not in HANDLED_FUNCTIONS:
            return NotImplemented
        return HANDLED_FUNCTIONS[func](*args, **kwargs)
    
def make_numericallabelled(x, label=""):
    """
    Helper function to cast anything into a NumericalLabelled object.
    """
    if isinstance(x, NumericalLabeled):
        return x
    else:
        return NumericalLabeled(x, label=label)
    
# Numpy functions            
# Override functions - used with __array_function__
def implements(np_function):
    def decorator(func):
        HANDLED_FUNCTIONS[np_function] = func
        return func
    return decorator    
    

@implements(np.random.normal)
def np_random_normal(loc=0.0, scale=1.0, **kwargs):
    # cast both loc and scale into Numericallabelled
    loc = make_numericallabelled(loc)
    scale = make_numericallabelled(scale)
    # check their label is "compatible"
    if not loc.label == scale.label:
        raise ValueError
    return NumericalLabeled(np.random.rand(loc=loc.value,
                                           scale=scale.value, **kwargs), 
                            loc.label+scale.label)



def main():
    # create two scalars
    half = NumericalLabeled(0.5, "half")
    # this raises a TypeError : 
    print(np.random.normal(loc=half))
    # TypeError: unsupported operand type(s) for *: 'NumericalLabeled' and 'int'

main()