如何将大小为1的NumPy一维数组转为标量？

我为此困惑了一小时。想实现一个函数，参数支持数字或一维数组，要求像ufunc那样：当输入是数字时返回标量（而非0维数组），输入是数组时返回同形状的数组。

当前函数实现如下：

def func(x: Real | np.ndarray, arr: np.ndarray):
    """Illustrate an actually more complicated  function"""
    return arr @ np.sin(arr[:, None] * x)

已知arr是一维数组，这里把它升为二维是为了避免元素级乘法的广播问题，但现在的问题是函数始终返回一维数组。不过这个实现能覆盖以下场景：

• x为标量且len(arr) == 1；
• x为标量且len(arr) > 1；
• x为数组且len(arr) == 1；
• x为数组且len(arr) >= 1，无论len(x)与len(arr)是否相等

我试过用@np.vectorize装饰器：

@np.vectorize
def func(x, arr):
    return arr @ np.sin(arr * x)

但还是始终返回一维数组，而且性能很差。我也考虑过functools.singledispatch，但会导致大量代码重复，还容易遗漏边界情况。

我自己想了一种解决方案：

def func(x, arr):
    res = arr @ np.sin(arr[:, None] * x)
    if len(res) == 1:
        return res.item()
    return res

但这类函数我有很多，这种写法不够Pythonic？于是我想写个装饰器来统一处理这个检查：

def give_me_a_scalar(f):
    @functools.wraps(f)
    def wrapper(*args, **kwargs):
        res = f(*args, **kwargs)
        if len(res) == 1:
            return res.item()
        return res
    return wrapper

这个装饰器似乎能满足需求，但我觉得应该有更简单的现成方案，是不是我遗漏了什么？

内容的提问来源于stack exchange，提问作者Aubergine