TensorFlow broadcasting of RaggedTensor

user_824276 0
Viewed 150

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
tensorflow2.0 tensorflow python
3 Answers

To achieve the desired effect, try the following:

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

Note that we only added tf.expand_dims() on the original y which is still tf.constant([[1], [2]])

The following desired result is obtained (tested on TF 2.6 version):

<tf.RaggedTensor [[[2, 0], [2]], [[0], [1, 2]]]>

As can be seen on the image below: Result of running the code

user_109854120

Does this work for you?

x = tf.ragged.constant([[[3, 1], [3]],
                        [[2], [3, 4]]])  # shape(2,None,None)

x = tf.RaggedTensor.from_uniform_row_length(x.values, uniform_row_length=2) # shape(2,2,None)

Output of x - y

<tf.RaggedTensor [[[2, 0], [1]],
 [[1], [1, 2]]]>

A preferable way would be to construct x as a tensor with ragged rank 1 from the beginning

x = tf.ragged.constant([[3, 1], [3],
                        [2], [3, 4]], ragged_rank=1)  # shape(4,None)
x = tf.RaggedTensor.from_uniform_row_length(x, uniform_row_length=2)  # shape(2,2,None)
user_89906440

Unfortunately, tf.RaggedTensors are a bit flaky.

tldr; Operations between tf.RaggedTensor and tf.Tensor do not always as expected. Sometimes converting tf.Tensor to a tf.RaggedTensor does the job.

The solution is as follows. X can remain as is,

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

which has a shape (X.shape) of [2, None, None].

Solution

Next we convert y to a tf.RaggedTensor as follows.

y = tf.ragged.constant(pylist=[[[1], [2]]], ragged_rank=1)

which has a shape (y.shape) of [1, None, 1]. And,

X-y gives,

<tf.RaggedTensor [[[2, 0], [1]], [[1], [1, 2]]]>

Caveat

For y a ragged_rank=2 will lead to the following error.

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

which I don't fully comprehend. The best way to avoid this is by having as many dimensions fixed as possible. And this will become more complex as the dimensionality of your tensors grow. For this, you can use RaggedTensor constructs such as from_uniform_row_length (More in for here and here)

user_16990750
Related