有没有可能使卷积层具有张量流概率

Question

在关于TensorFlow概率中概率层回归的教程中，它建立了以下模型：

model = tfk.Sequential([
  tf.keras.layers.Dense(1 + 1),
  tfp.layers.DistributionLambda(
      lambda t: tfd.Normal(loc=t[..., :1],
                           scale=1e-3 + tf.math.softplus(0.05 * t[..., 1:]))),
])

我的问题是，卷积层也可以这样做吗？例如，在3D分割中，输出层具有形状（3、128、128、128）。如何使最后一层（或登录）具有上述概率？如果我展平最后一层，那么维度是：3*128*128*128，这将导致资源耗尽错误