Tensorflow tf.map_fn 在不规则张量上失败，类型为“RaggedTensor”的对象没有 len

Question

This Tensorflow doc gives this example of using tf.map_fn on ragged tensors which works for Tensorflow 2.4.1 and above:

digits = tf.ragged.constant([[3, 1, 4, 1], [], [5, 9, 2], [6], []])
print(tf.map_fn(tf.math.square, digits))

However the following example results in error "object of type 'RaggedTensor' has no len" when run in Tensorflow 2.4.1 or Tensorflow 2.5:

import tensorflow as tf

X=tf.ragged.constant([[1.,2.],[3.,4.,5.]], dtype=tf.float32)

@tf.function
def powerX(i):
    global X
    return X**i

Y = tf.map_fn(powerX, tf.range(3, dtype=tf.float32))

Is there a way to make this work? I don't understand the error being thrown. In general I am trying to get full parallelism by mapping a user defined function which has only Tensorflow operations over a ragged tensor with results being ragged tensors.

Answer 1

tf.map_fn requires an output signature. I'm not sure why it can't infer this from the input but that's a question for the tensorflow people. The following code will work for you.

import tensorflow as tf

X=tf.ragged.constant([[1.,2.],[3.,4.,5.]], dtype=tf.float32)

@tf.function
def powerX(i):
    global X
    return X**i
signature = tf.type_spec_from_value(powerX(X))
Y = tf.map_fn(powerX, tf.range(3, dtype=tf.float32),fn_output_signature=signature)