How can I fill a ragged tensor with zeros to make a square tensor

Question

I want to transform A into B:

A = 
  [1, 1, 1]
  [2]
  [3, 3, 3, 3]
  [4, 4]

B =
  [
   [0, 1, 1, 1]
   [0, 0, 0, 2]
   [3, 3, 3, 3]
   [0, 0, 4, 4]
  ]

Input :
- a list of lists

Output :
- a single matrix or tensor
- right aligned
- Left fill with 0's

Answer 1

In specific case where values are same per first dimension as are in your example. You can use:

digits = tf.ragged.constant([[1., 1., 1.],[2.],[3., 3., 3., 3.],[4., 4.]])
padded = digits.to_tensor(0.)
final_tensor = tf.reverse(padded, [-1])

output:

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

Answer 2

A is a list, so we can for-loop and use tf.pad and tf.stack to get the output tensor.

max_len = max(len(e) for e in A)
res = tf.stack([tf.pad(e, [[max_len - len(e),0]]) for e in A], axis=0)
# array([[0, 1, 1, 1],
#        [0, 0, 0, 2],
#        [3, 3, 3, 3],
#        [0, 0, 4, 4]], dtype=int32)