I always wonder how to design a vectorizable function when the sizes of the input and output arrays are different.
In other words, I am always unsure whether to shift the “core dimension” to the back
def polar_coordinates(r, theta):
xp = array_api_compat(r, theta)
return xp.stack([r * xp.cos(theta), r * xp.sin(theta)], axis=-1)
or to the front
def polar_coordinates(r, theta):
xp = array_api_compat(r, theta)
return xp.stack([r * xp.cos(theta), r * xp.sin(theta)], axis=0)
Is there any plans to add recommendations for this to array API? For reference