Add batch_gather that assumes a series of leading (dense) dimensions.

This operation computes:

 result[i1, ..., in] = params[i1, ..., in-1, indices[i1, ..., in]]

That is, it assumes that `params` and `indices` have n-1 leading dense dimensions, followed by a sparse dimensions over which we want to gather.

PiperOrigin-RevId: 208836509