Wednesday, 31 August 2022

TensorFlow broadcasting of RaggedTensor

How do I subtract a tensor from ragged tensor?

Example:

import tensorflow as tf    # TensorFlow 2.6

X = tf.ragged.constant([[[3, 1], [3]],
                        [[2], [3, 4]]], ragged_rank=2)
y = tf.constant([[1], [2]])
X-y

Expected result:

[[[2, 0], [1]],
 [[1], [1, 2]]]

However, it returns an error:

tensorflow.python.framework.errors_impl.InvalidArgumentError: Expected 'tf.Tensor(False, shape=(), dtype=bool)' to be true. Summarized data: b'Unable to broadcast: dimension size mismatch in dimension'
1
b'lengths='
2
b'dim_size='
2, 2

I know I can do it row-by-row:

result = []
if X.shape[0] is not None:  # Placeholders have None everywhere -> range(None) raises an exception -> this condition
    for row in range(X.shape[0]):
        result.append(X[row] - y)
    result = tf.stack(result)

However, this works only in the eager mode - in graph mode, I get:

ValueError: No gradients provided for any variable

because the code gets executed only conditionally...

What works is to hard-code the count of rows:

for row in range(2):
    result.append(X[row] - y)
result = tf.stack(result)

But that doesn't generalize well.

I also know I can write:

X - tf.expand_dims(y, axis=1)

But that returns a result for "transposed" y:

[[[2, 0], [2]],
 [[0], [1, 2]]]

I also know I can also use:

def subtract(x):
    return x - y

tf.map_fn(subtract, X)

But when using the result in graph mode, I get:

ValueError: Unable to broadcast: unknown rank


from TensorFlow broadcasting of RaggedTensor

No comments:

Post a Comment