How to apply vmap on a heterogeneous tensor

I believe those tensors I don’t need a jagged layout.

Yep

But if I have a bunch of tensors with shape [(n, ), (m, )...] , then I need jagged?

Jagged layout only supports a single jagged dimension, e.g. among the tensors that you have, only a single dimension can vary. So if you have a bunch of 2-D tensors where n/m can both vary, then jagged layout would not work.

Is that jagged mean the dimension in the middle is irregular?

Your jagged dimension cannot be the very first one because there needs to be a preceeding batch dimension.
But you can have a 2D nested tensor, e.g. with shape= [3, [1, 2, 3]] where the varying dimension is the very last one.