chore: use xp.take_along_axis
if Array API version >=2024.12
#10253
Loading
xp.take_along_axis
if Array API version >=2024.12
#10253