Repository navigation
Preferences regarding "core dimension" #966
Description
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
-
Numpy's gufunc puts "core dimension" to the back https://numpy.org/doc/stable/reference/c-api/generalized-ufuncs.html
-
Recently added
scipy.special.sph_harm_y_allputs "core dimension" to the front https://docs.scipy.org/doc/scipy/reference/generated/scipy.special.sph_harm_y_all.html -
Since
broadcastingadjusts shapes to the back, by putting "core dimension" to the back one can more easily interact with newly added feature (dimension), and vice versa. -
In terms of calculation speed, this should be related to C-style and Fortran-style indexing.
Activity
rgommers commentedon Aug 22, 2025
I don't think so. Given that we're rarely designing new APIs and usually adopting what is already common across array libraries, we haven't had to give this one any thought.
Both broadcasting and batching (
linalgAPIs, deep learning operators) grow new dimensions at the front, so that seems to be the prevalent pattern - gufuncs are niche in comparison. However, as long as you document it well, either choice seems reasonable and has precedent.34j commentedon Aug 22, 2025
I see, I'm starting to feel like either would be fine.
Regarding the speed, considering that many libraries use C-style indexing, is it okay to say that ordinary vectorizable array API compatible code would be faster if it grows "new dimensions at the front"?
ev-br commentedon Sep 15, 2025
As a data point,
scipy.interpolatetypically adds batch dimensions to the end. This did cause some mild inconveniences, mostly because it's not what usual broadcasting does.