我有一个自定义的tensorflow层，它可以很好地生成输出，但在与Keras函数模型API一起使用时会抛出错误。

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)

这将生成以下输出：

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

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

但是当我在Keras函数API中使用它时，它不起作用。

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

它给出以下错误：

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.

有人知道为什么会发生这种情况，以及如何修复吗？