Как перебирать "некоторые" размеры массива numpy?

У меня есть трехмерный массив формы (2,2,2). Я хотел бы думать об этом как о двумерной матрице с одномерными массивами в качестве записей.

Я хотел бы применить функцию к каждой записи (то есть к каждому 1-мерному массиву) в моей матрице. Я знаю, что могу векторизовать свою функцию, чтобы применить ее к каждому числу в моем массиве. Я также знаю, что могу применить функцию по одной оси. Однако мне не удалось применить функции по двум осям.

Вот мое последнее испытание:

import numpy as np

def sqrtSum(a, b):
    return np.sqrt(a+b)

def sqrtSumWrapper(row):
    return np.array([sqrtSum(x[0], x[1]) for x in row])

z = np.array([[[1,2],[3,4]],[[5,6],[7,8]]])
np.apply_along_axis(sqrtSumWrapper, 1, z)

В приведенном выше примере моим желаемым результатом был бы массив формы (2,2) с записями ((sqrt (3), sqrt (7)), (sqrt (11), sqrt (15)). Когда я запускаю код выше, я получаю ошибку неверный индекс для скалярной переменной..

Я думаю, что мне не хватает важного аспекта того, как работает apply_along_axis, и я был бы благодарен за ваш вклад в то, как исправить код.

Обновлено: Пока что ответы сосредоточены на изменении входной функции sqrtSum. Эта функция - всего лишь пример. Меня интересует общий ответ для произвольной функции ввода, которая принимает n входных параметров и возвращает скаляр.

Решение: Решение на удивление простое (позор мне, что я его не видел)

import numpy as np
def sqrtSum(a, b):
    return np.sqrt(a+b)

def sqrtSumWrapper(x):
    return sqrtSum(x[0], x[1])

z = np.array([[[1,2],[3,4]],[[5,6],[7,8]]])
np.apply_along_axis(sqrtSumWrapper, 2, z)

Спасибо всем, кто ответил.