为什么TensorFlow 2自定义Bijector/flow生成的样本数量不符合预期

核心是自定义Bijector的_forward方法索引逻辑错误，没有兼容batch维度。
当调用nf.sample(3)时，底层基础分布首先会返回形状为(3, 2)的样本张量（3个batch，每个是2维事件）。原有写法z[0], z[1]是直接取张量的第0和第1个batch的全部元素，每个的形状是(2,)，拼接后返回的张量形状为(2, 2)，这就是仅得到2个样本的原因。

需要做两处调整：

• 调整_forward方法的取值逻辑，从最后一个维度（事件维度）拆分元素，兼容任意batch维度
• 修正forward_min_event_ndims参数，该Bijector操作的是1维的事件向量，所以参数应设为1，避免后续计算对数概率时出现维度匹配错误

修改后的完整Bijector代码如下：

class F3Bijector( tfb.Bijector ):
    """ 
    Represents bijector from [ z1, z2 ] to [ z1, z2 + a * z1 ** 2 ]
    """
    def __init__( self, a, **kwargs ):
        tfb.Bijector.__init__( self, forward_min_event_ndims = 1, **kwargs )
        self.a = a

    def _forward( self, z ):
        z1 = z[..., 0]
        z2 = z[..., 1]
        return tf.stack([ z1, z2 + self.a * z1 ** 2 ], axis=-1)   

修改后调用nf.sample(3)就会正确返回形状为(3, 2)的3个样本。

内容的提问来源于stack exchange，提问作者Mark Lavin