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

Thanks @carlosgmartin, I don't know if I think it's very useful, but I think that places that accept axis and multiple make clearly sense and is simple, then it's good to just do it.

@seberg In addition to #29240 (comment), this is needed by reductions like mean and var, which need to know the size of the axis or axes along which the reduction takes place.

Thus the endpoint of the chain mean > _mean > _count_reduce_items does something similar.

this is needed by reductions like mean and var, which need to know the size of the axis or axes along which the reduction takes place.

I don't like np.size to begin with much, since it does coercion to ndarray first potentially. That function is needed because of the masked path, the other path could be simplified, but whether it's better off calling np.size or just simplifying the logic here, dunno.

The non-masked path could perhaps reuse the new np.size, whereas the masked path could stick to where.sum(axis=axis, keepdims=keepdims) or something similar.