средний массив numpy, но сохраняющий форму

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

a = np.array([[[0.1, 0.2, 0.3], [0.2, 0.3, 0.4]],
              [[0.4, 0.4, 0.4], [0.7, 0.6, 0.8]]], np.float32)

b = np.average(a, axis=2)
# b = [[0.2, 0.3],
#      [0.4, 0.7]]

Требуемый результат:

# b = [[[0.2, 0.2, 0.2], [0.3, 0.3, 0.3]],
#      [[0.4, 0.4, 0.4], [0.7, 0.7, 0.7]]]

Можете ли вы сделать это элегантно или мне просто нужно перебрать массив в Python (который будет намного медленнее по сравнению с мощной функцией Numpy).

Можете ли вы установить аргумент Dtype для функции np.mean в одномерный массив?

Спасибо.