padded_batch: pad 2D tensors with a 1D tensor

Question

I have a dataset whose elements are 2D tensors. These elements have shape (None, 2). An example dataset contains these two elements:

[[1 2]
 [3 4]]

[[ 5  6]
 [ 7