To add to piotr’s answer, if your tensors on differ in shape only on a single dimension you can try NestedTensors. How to apply vmap on a heterogeneous tensor - #2 by soulitzer
Internally nested tensors represent ragged data using a packed representation.
You might be able to get more coverage on certain operators compared to foreach.