迭代numpy数组列的循环（维度为1D或2D）

Question

我正在尝试编写一个函数，它获取一个numpy数组，INPUT，并将它的列逐个传递给另一个函数。 INPUT数组是1D或2D（不是更多），第二个函数需要1D数组作为参数。 （LEN（param.shape）== 1）

我已经阅读了一个类似的线程，其中OP想要对所有列进行求和并检查其他条件......这可能需要另一个答案。

伪代码中的所需操作：

def func(INPUT,a,b,...)
    for column in INPUT: #whether be a 1D or 2D
        result = another_func(column,...)

试过这个：问题是如何不在func中检查INPUT数组的维度：

if(len(INPUT.shape)==1):
    another_func(INPUT,....)
elif(len(INPUT.shape)==2):
    for c in range(INPUT.shape[1]):
        another_func(INPUT[:,c])