Tensorflow layer working outside of model but not inside

January 11, 2023

I have a custom tensorflow layer which works fine by generating an output but it throws an error when used with the Keras functional model API. Here is the code:

import numpy as np
import tensorflow as tf
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input

# ------ Custom Layer -----------
class CustomLayer(tf.keras.layers.Layer):
  def __init__(self):
    super(CustomLayer, self).__init__()

  def split_heads(self, x):
    
    batch_size = x.shape[0]
    split_inputs = tf.reshape(x, (batch_size, -1, 3, 1))
    
    return split_inputs
  
  def call(self, q):

    qs = self.split_heads(q)

    return qs

# ------ Testing Layer with sample data --------
x = np.random.rand(1,2,3)
values_emb = CustomLayer()(x)
print(values_emb)

This generates the following output:

tf.Tensor(
[[[[0.7148978 ]
   [0.3997009 ]
   [0.11451813]]

  [[0.69927174]
   [0.71329576]
   [0.6588452 ]]]], shape=(1, 2, 3, 1), dtype=float32)

But when I use it in the Keras functional API it doesn’t work. Here is the code:

x = Input(shape=(2,3))
values_emb = CustomLayer()(x)
model = Model(x, values_emb)
model.summary()

It gives this error:

TypeError: Failed to convert elements of (None, -1, 3, 1) to Tensor. Consider casting elements to a supported type. See https://www.tensorflow.org/api_docs/python/tf/dtypes for supported TF dtypes.

Does anyone know why this happens and how it can be fixed?