I don't mind putting this in, but the discussion did not actually point to why this is useful.

It would have been nice to make axis a kwonly argument, which unfortunately is not a backward compatible change. I would need to reach for the documentation if I ever saw np.size(a, 1), and even moreso if I saw np.size(a, (1, 1)). Perhaps at least for the examples we should be using the kwarg?

@mattip This is useful for measuring the size of a multidimensional array with semantically-meaningful axes in a specific, nontrivial way.

For example, I can get the pixel size of an image in HWC format via jnp.size(array, (0, 1)), or in CHW format via jnp.size(array, (1, 2)). I can get the voxel size of a volumetric image via jnp.size(array, (0, 1, 2)). These might have leading batch dimensions, so I can adjust the tuple accordingly. There might also be a temporal dimension, for videos or spatio-temporal data/simulations. And so on.

Thanks. Personally i use dtype.itemsize * factor, but I guess using size works too