TypeError: Could not build a TypeSpec for <KerasTensor when using tf.map_fn and keras functional model

user_7381780 0
Viewed 2155

When I attempt to use tf.map_fn in the definition of a keras Functional model, I get the error:

TypeError: Could not build a TypeSpec for <KerasTensor: ...

e.g. this simple model will trigger that error in tf-nightly 2.5.0 :

import tensorflow as tf

from tensorflow.keras.layers import Input
from tensorflow.keras.models import Model

x = Input(shape=(10,))

y = tf.map_fn(lambda x : x * 2, x, fn_output_signature=tf.float32)

model = Model(inputs=x, outputs=y)

Whereas replacing the call to tf.map_fn to calls to other tensorflow operations works fine.

keras tensorflow python
2 Answers

This seems to result from tf.map_fn being unable to determine the TypeSpec of the input tensor when the input is a keras sympolic Input.

Several times now I've run into problems with tensorflow ops and Keras symbolic Input tensors. Wrapping the offending code in a custom layer seems to generally fix it.

For example, replacing the code above with this will execute successfully:

import tensorflow as tf

from tensorflow.keras.layers import Input, Layer
from tensorflow.keras.models import Model

x = Input(shape=(10,))

class MapLayer(Layer):

    def call(self, input):
        return tf.map_fn(lambda x : x * 2, input, fn_output_signature=tf.float32)

y = MapLayer()(x)

model = Model(inputs=x, outputs=y)
user_73817800

Why don't you use a tf.keras.layers.Lambda layer instead?

import tensorflow as tf

from tensorflow.keras.layers import Input
from tensorflow.keras.models import Model

x = Input(shape=(10,))

y = tf.keras.layers.Lambda(lambda x : x * 2)(x)

model = Model(inputs=x, outputs=y)
user_109083750
Related