Tensorflow Keras Lambda层，如何使用多个参数

我有这个功能：

def sampling(x):
    zeros = x*0
    samples = tf.random.categorical(tf.math.log(x), 1)
    samples = tf.squeeze(tf.one_hot(samples, depth=2), axis=1)
    return zeros+samples

我从这个层调用：

x = layers.Lambda(sampling, name="lambda")(x)

但我需要更改采样函数中的深度变量，因此我需要如下内容：

def sampling(x, depth):

但是，我怎样才能使它与Lambda层一起工作呢

---

非常感谢

在lambda层中使用lambda函数

def sampling(x, depth):
    zeros = x*0
    samples = tf.random.categorical(tf.math.log(x), 1)
    samples = tf.squeeze(tf.one_hot(samples, depth=depth), axis=1)
    return zeros+samples

用法：

Lambda(lambda t: sampling(t, depth=3), name="lambda")(x)

---

成功了，谢谢！