How do I implement itertools.permutations in tensorflow

I'm trying to find a way to generate a tensor of permutations using tensorflow.

Basically, I'd like to have a function like the following:

out = tf_permutations(tf.range(3))

Where the output would be:

<tf.Tensor: shape=(6, 3), dtype=int32, numpy=
array([[0, 1, 2],
       [0, 2, 1],
       [1, 0, 2],
       [1, 2, 0],
       [2, 0, 1],
       [2, 1, 0]], dtype=int32)>