Use gather as the batched form of dynamic_slice.
Use gather as the batched form of dynamic_slice.

Avoids building an unrolled slice/concatenate structure that is linear in the size of the batch dimension.

tensorflow/probability,google/jax,google/jax,tensorflow/probability,google/jax,google/jax