如何從 tensorflow 中的 2d boolean 掩碼沿第 0 軸找到“下一個”值

Question

假設我有一個 boolean 張量，它告訴我給定坐標處的值是否“感興趣”：

is_value_of_interest = np.array(
    [[0, 0, 0],
    [0, 1, 1],
    [1, 0, 0],
    [0, 0, 0],
    [0, 1, 1]])

is_value_of_interest_tf = tf.constant(is_value_of_interest)

現在我有另一個包含實際值本身的數組/張量：

values = np.random.rand(5, 3)
values_tf = tf.constant(values)

我想要做的是構建一個張量，它將沿第 0 軸返回下一個感興趣的值。 所以在 numpy/pandas 中，這將是：

values_of_interest = np.where(is_value_of_interest, values, np.nan)
df = pandas.DataFrame(values_of_interest).bfill()

結果：

我如何構建張量操作來實現相同的結果，例如(df.values == my_tensor.numpy()).all() ？

Answer 1

由於缺少與 pandas 的 ffill/bfill 方法等效的方法，因此這是非常重要的，因此首先在沿第 0 維的 2d 輸入的特定情況下為 tensorflow 實施ffill function。

我對結果不滿意，因為它又長又亂，所以如果有人可以改進它，請隨時提出更好的答案：

def tf_ffill(data: tf.Tensor) -> tf.Tensor:
    """
    2d forward-fill along 0th dimension
    """

    is_value_of_interest = ~tf.math.is_nan(data)
    grid = tf.meshgrid(tf.range(data.shape[0]), tf.range(data.shape[1]), indexing="ij")
    nan_sentinel = tf.maximum(data.shape[0], data.shape[1])

    next_value_of_interest_row_index = (
        tf.cumsum(tf.cast(is_value_of_interest, tf.int32), axis=0) - 1
    )
    defined_values_mask = tf.where(next_value_of_interest_row_index == -1, False, True)

    next_value_of_interest_row_index = tf.where(
        defined_values_mask, next_value_of_interest_row_index, nan_sentinel
    )
    next_value_of_interest_col_index = tf.where(
        defined_values_mask, tf.cast(grid[1], tf.int32), nan_sentinel
    )

    # this is a rank N+1 tensor with the last 2 dimensions representing the x/y
    # coordinates of the index of the next point of interest The only quirk
    # here is that the column is the 2nd last dimension and the row is the
    # last. This is to align it with the representation produced by the
    # boolean_mask op.
    values_of_interest_index_lookup = tf.concat(
        [
            tf.expand_dims(next_value_of_interest_col_index, axis=2),
            tf.expand_dims(next_value_of_interest_row_index, axis=2),
        ],
        2,
    )
    values_of_interest_indices = tf.ragged.boolean_mask(
        tf.transpose(grid[0]), tf.transpose(is_value_of_interest)
    )

    mapped_indices = tf.gather_nd(
        values_of_interest_indices,
        tf.where(
            values_of_interest_index_lookup == nan_sentinel,
            0,
            values_of_interest_index_lookup,
        ),
    )
    mapped_indices = tf.where(defined_values_mask, mapped_indices, nan_sentinel)

    res = tf.experimental.numpy.take_along_axis(
        data, tf.where(mapped_indices == nan_sentinel, 0, mapped_indices), axis=0
    )
    return tf.where(defined_values_mask, res, np.nan)


def tf_bfill(data: tf.Tensor) -> tf.Tensor:
    return tf.reverse(tf_ffill(tf.reverse(data, axis=[0])), axis=[0])

一旦你有tf_bfill得到答案是微不足道的：

data = tf.where(is_value_of_interest, values_tf, np.nan)
tf_bfill(data)

整個問題可以改寫為“你如何在 tensorflow 中進行正向填充？ ”