tf.vectorized_map does not concat variable length tensors (InvalidArgumentError: PartialTensorShape: Incompatible shapes during merge)

The problem is you are passing a tensor to tf.ones and tf.zeros instead of a shape. For example, if you pass the tensor a to tf.ones, it will be interpreted as the shape resulting in a tensor with the shape (5, 4, 3, 2). That is probably not what you want. Try something like this:

import tensorflow as tf 

def test_fn(inputs):
  a, b = inputs
  out = tf.stack([tf.ones_like(a), tf.zeros_like(b)], 0)
  return out

a = tf.constant([5,4,3,2])
b = tf.constant([5,6,7,8])
x_a = tf.vectorized_map(test_fn,(a,b))
x_a = tf.transpose(x_a)
print(x_a)

tf.Tensor(
[[1 1 1 1]
 [0 0 0 0]], shape=(2, 4), dtype=int32)

Note that you have to use tf.stack instead of tf.concat because TF does not currently support scalar concatenation when using tf.vectorized_map. Check out the limitations of tf.vectorized_map here.