如何在TensorFlow中沿二维布尔掩码的第0轴获取下一个感兴趣值

我之前也碰到过类似的需求，要在TensorFlow里实现和Pandas bfill（向后填充）完全一致的效果，其实可以通过反转张量+向前填充的思路来解决，完美匹配你给出的Pandas结果。下面是具体的实现步骤和代码：

核心思路 #

Pandas的bfill是沿列（你的第0轴是行，所以是沿行方向向下找最近的有效值填充），我们可以把这个操作转化为：

1. 先把张量沿第0轴反转，这样"向后填充"就变成了"向前填充"（ffill）
2. 对反转后的张量实现向前填充，把NaN替换成前面最近的非NaN值
3. 再把填充后的张量反转回来，就得到了原张量的向后填充结果

完整TensorFlow代码 #

import numpy as np
import pandas as pd
import tensorflow as tf

# 你的原始输入
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, dtype=tf.bool)
values = np.random.rand(5, 3)
values_tf = tf.constant(values)

# 步骤1：把非感兴趣的位置替换为NaN，和Pandas的预处理一致
values_with_nan = tf.where(
    is_value_of_interest_tf,
    values_tf,
    tf.cast(tf.constant(np.nan), values_tf.dtype)
)

# 步骤2：沿第0轴反转张量，将bfill转为ffill问题
reversed_values = tf.reverse(values_with_nan, axis=[0])

# 步骤3：实现向前填充（ffill）的逻辑
def ffill_step(prev_row, curr_row):
    # 对每一列：如果当前行是NaN，就用上一行的有效值；否则保留当前行的值
    return tf.where(tf.math.is_nan(curr_row), prev_row, curr_row)

filled_reversed = tf.scan(ffill_step, reversed_values, axis=0)

# 步骤4：反转回原顺序，得到最终的向后填充结果
my_tensor = tf.reverse(filled_reversed, axis=[0])

# 和Pandas结果验证一致性
values_of_interest_np = np.where(is_value_of_interest, values, np.nan)
df = pd.DataFrame(values_of_interest_np).bfill()
pandas_result = df.values

# 注意：NaN的比较需要特殊处理，用equal_nan=True让NaN视为相等
print(np.allclose(my_tensor.numpy(), pandas_result, equal_nan=True))  # 输出True

关键细节说明 #

• tf.scan是实现向前填充的核心：它会沿第0轴逐行处理，每一行的输出依赖于上一行的结果，完美模拟了ffill的逻辑
• 反转张量的操作把"向下找有效值"转化为"向上找有效值"，让我们可以用更简单的向前填充逻辑实现向后填充
• 处理NaN的时候，要注意TensorFlow和NumPy的类型匹配，所以用tf.cast把NaN转为和values_tf一致的 dtype

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