在Keras层中使用自定义操作，并使张量可训练

import tensorflow as tf
from tensorflow.keras import layers

class CustomLayer(layers.Layer):
    def __init__(self, units=32):
        super(CustomLayer, self).__init__()
        self.units = units

    def build(self, input_shape):
        # 定义可训练的参数
        self.w = self.add_weight(shape=(input_shape[-1], self.units),
                                 initializer='random_normal',
                                 trainable=True)
        self.b = self.add_weight(shape=(self.units,),
                                 initializer='zeros',
                                 trainable=True)

    def call(self, inputs):
        # 自定义操作
        return tf.matmul(inputs, self.w) + self.b

# 创建一个包含自定义层的模型
model = tf.keras.Sequential([
    CustomLayer(units=64),
    layers.Activation('relu'),
    layers.Dense(units=10)
])

# 编译模型
model.compile(optimizer=tf.keras.optimizers.Adam(),
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])

# 训练模型
model.fit(x_train, y_train, epochs=10, batch_size=32)