Why does dot_general move batch dimensions?

[otoomey]

I am currently trying to understand how dot_general works. I noticed something that strikes me as a bit strange:

#                b  -  j  b
lhs = jnp.zeros((2, 5, 2, 3))
#                b  j  -  b
rhs = jnp.zeros((2, 2, 7, 3))

lax.dot_general(lhs, rhs, [([2], [1]), ([0, 3], [0, 3])]).shape
# -> (2, 3, 5, 7)

I would have expected the result to have a shape of (2, 5, 7, 3), i.e. the final batch dimension is still in the same place. From a performance standpoint it may well make sense to prepend all the batch dimensions, but surely jax could just put the dimension "back" where it belongs after computation. Is there a reason why it works this way?

[jakevdp]

JAX follows StableHLO/OpenXLA semantics for DotGeneral, described here: https://openxla.org/xla/operation_semantics#dotgeneral. In particular:

It follows that the resulting dimension number starts with the batch dimension, then the lhs non-contracting/non-batch dimension, and finally the rhs non-contracting/non-batch dimension.

As for why it does this: it's a straightforward convention. One can imagine other straightforward conventions, but this is the one that was chosen by the designers of XLA. I suspect the reason is that in ML computations over C-ordered arrays, batch dimensions are typically laid out in the leading dimension.