Custom Keras Layer fails with "TypeError: Could not build a TypeSpec for KerasTensor"-CodePudding

I reduced my problem to this very basic example:

from tensorflow import keras
from tensorflow.keras import layers

class MyLayer(layers.Layer):
    def __init__(self, z):
        super(MyLayer, self).__init__()
        self.z = z

    def call(self, inputs):
        return inputs * self.z
      
inputs = keras.Input(shape=(2,))
layer = MyLayer(inputs)
z = layer(inputs)

model = keras.Model(inputs=inputs, outputs=z)

Which fails with

TypeError: Could not build a TypeSpec for <KerasTensor: shape=(None, 2) dtype=float32 (created by layer 'tf.math.multiply_17')> with type KerasTensor

Do you have any idea where is the problem?

CodePudding user response：

You cannot multiply a Keras symbolic tensor with a Tensorflow tensor in graph mode. There are a few ways to do what you want IIUC, here is one option assuming you want to pass two tensors to MyLayer:

from tensorflow import keras
from tensorflow.keras import layers

class MyLayer(layers.Layer):
    def __init__(self):
        super(MyLayer, self).__init__()

    def call(self, inputs):
        x, y = inputs
        return x * y
      
inputs = keras.Input(shape=(2,))
layer = MyLayer()
z = layer([inputs, inputs])

model = keras.Model(inputs=inputs, outputs=z)