Python 如何实例化tfp.layers.DistributionLambda中包含分布的变量

我正在尝试使用

tfp.layers.DistributionLambda

构建一个

tf.keras.Sequential

模型。我正在遵循，但希望用包含

tfd.TransformedDistribution

的变量替换

tfd.Normal

，并使用

RealNVP

bijector

import tensorflow as tf
import tensorflow_probability as tfp


model = tf.keras.Sequential((
    tf.keras.layers.Lambda(lambda x: tf.shape(x)[-1]),
    tfp.layers.DistributionLambda(lambda t: (
        tfp.distributions.TransformedDistribution(
            distribution=(
                tfp.distributions.MultivariateNormalDiag(loc=tf.zeros(t))),
            bijector=tfp.bijectors.RealNVP(
                num_masked=2,
                shift_and_log_scale_fn=tfp.bijectors.real_nvp_default_template(
                    hidden_layers=[32, 32]))))),
))


x = tf.random.uniform((5, 3))
distribution = model(x)

但是，这在以下情况下失败：

该层无法安全地确保跨多个应用程序正确重用变量 因此，出于安全考虑，这种行为是不允许的。兰姆达 层不适合于有状态计算；相反，写一个 子类层是使用 变量

请注意，

RealNVP

bijector的变量必须在bijector内初始化，这与中的

Normal

分布变量不同，后者是在

Sequential

的顶层创建的

我想知道是否有一种方法可以将

DistributionLambda

用于必须在分布中创建变量的设置？如果是这样，处理

DistributionLambda

层内变量的正确方法是什么？如果不是，那么建议用什么方法构建这样的模型