python - numpy apply along axis with error "ValueError: could not broadcast input array from shape (2) into shape (1)" -

here's example of numpy array, , apply function mapping each rows of matrix test.

test = np.array([[0, .1, .9], [.1, .9, .8], [.8, .6, .1]]) test2 = np.array(['a','b','c'])  def mapping(x):     return test2[np.where(x > .7)].tolist() 

this works

mapping(test[0]), mapping(test[1]), mapping(test[2]) 

correct result: (['c'], ['b', 'c'], ['a','b'])

but doesn't, , spits out error.

np.apply_along_axis(mapping, 1, test) 

i don't understand why case. please help.

from apply docs:

the output array. shape of `outarr` identical shape of `arr`, except along `axis` dimension. axis removed, , replaced new dimensions equal shape of return value of `func1d`. if `func1d` returns scalar `outarr` have 1 fewer dimensions `arr`. 

what shape of return value of mapping? apply... tries guess doing calculation first item. in case ['c']. attempts return (3,1) array, , runs problems when 2nd row returns values.

the best way apply function rows of test is:

in [240]: [mapping(x) x in test] out[240]: [['c'], ['b', 'c'], ['a']]  in [246]: np.apply_along_axis(mapping, 1, test[[0,2],:]) out[246]:  array([['c'],        ['a']],       dtype='<u1') 

even works, apply_along_axis not improve speed - in fact it's worse

in [249]: timeit [mapping(x) x in test] 20.4 µs ± 33 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each) in [250]: timeit np.array([mapping(x) x in test]) 25.1 µs ± 192 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each) in [251]: timeit np.apply_along_axis(mapping, 1, test[[0,2,2],:]) 146 µs ± 194 ns per loop (mean ± std. dev. of 7 runs, 10000 loops e