Search code examples
pythontensorflowpadding

How to pad all tensors to accommodate all sizes


How do I pad two tensors of the same shape such that they both have the same size? I currently receive an error for the following code (I can't use for loops in graph mode):

def match_size(self, x, y):
    d = tf.maximum(tf.subtract(y.shape, x.shape), 0)
    x = tf.pad(x, [[0, i] for i in d])
    d = tf.maximum(tf.subtract(x.shape, y.shape), 0)
    y = tf.pad(y, [[0, i] for i in d])
    return x, y

This code will be run within a Keras model's call method due to the fact that x and y tensors will vary in feature size (last dim in shape (batch, horizon, feature)) throughout various stages of execution (i.e., I can't decide ahead of time during build what the sizes/shapes will be).

The following are the intended input/output examples:

x = (10, 4, 4), y = (10, 4, 2) ~> x = (10, 4, 4), y = (10, 4, 4)
x = (10, 4, 2), y = (10, 4, 4) ~> x = (10, 4, 4), y = (10, 4, 4)
...it should also work for all dimensions:
x = (10, 3, 2), y = (10, 4, 1) ~> x = (10, 4, 2), y = (10, 4, 2)


Solution

  • Using ragged.stack:

    def match_size(x, y):
       new_axis = tf.argsort(tf.abs(tf.constant(x.shape) - tf.constant(y.shape)), direction='DESCENDING')
       x_y = tf.ragged.stack([tf.transpose(x, new_axis),tf.transpose(y, new_axis)])
       x,y = x_y.to_tensor(0.)
       return tf.transpose(x, tf.argsort(new_axis)), tf.transpose(y, tf.argsort(new_axis))