tf.布尔掩模(2d,2d)给出1d结果

问题

tf.boolean_mask() 看起来相对简单:它从不匹配条件的张量中移除值。如果遮罩的尺寸与目标张量相同,则条件按元素方式工作。

用 ret = tf.boolean_mask(src, mask) ,我发现输出尺寸与输入尺寸不匹配。

src:    Tensor("mul_3:0", shape=(?,?), dtype=int32)
mask:   Tensor("Cast_1:0", shape=(?,?), dtype=int32)
ret:    Tensor("boolean_mask/Gather:0", shape=(?,), dtype=int32)

注意 (?,) 行为不同于 (?,?) 尽管我不知道为什么或如何。

代码

    src = [1, 2, 3, 4, 5, 6, 7, 8, 9]
    src = tf.ones([tf.shape(src)[0], 1], tf.int32) * src

    matrix = tf.ones_like(src, tf.int32)
    matrix = tf.matrix_band_part(matrix, 3, 3) # number for mask is arbitrary
    mask = tf.cast(matrix, tf.bool)

    ret = tf.boolean_mask(tensor=src, mask=mask)

怎么了?