Question

I have an function that maps an ndarray of shape (3) to a float, and I have an ndarray of shape (…,3). What’s the best way to map that function over that array to get an array of shape (…)?

Thanks.

Answer 1

You want numpy.apply_along_axis.

def f(a):
     return a[0] + a[1] + a[2]
mm = numpy.random.randn(5, 3)
numpy.apply_along_axis(f, 1, mm)

output: array([-1.75875289, -0.34689792, 0.66092486, -0.21626001, -0.14125476])