def apply2d_along_first(func2d, arr3d):
    a, n, m = arr3d.shape

def func1d(arr1d):
    return func2d(arr1d.reshape((n, m)))
arr2d = arr3d.reshape((a, n * m))
return apply_along_axis(func1d, -1, arr2d)
A, N, M = 3, 4, 5
data = arange(A * N * M).reshape((A, N, M))
print(data)
print(apply2d_along_first(trace, data))
Функции отображения двумерных массивов numpy