tf.boolean_mask (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)

В чем проблема?