TensorFlow中张量遇值1时修改后续i个元素的高效实现方法

现有两个尺寸完全一致的张量：

a = [1, 2, 3, 4, 5, 10, 11, 12, 13, 20, 21, 22, 23, 24, 25, 26, 27, 28]
b = [0, 1, 1, 1, 1,  1,  1,  1,  0,  1,  1,  1,  1,  1,  0,  0,  0,  1]

张量a包含3个按连续数值划分的独立区域：

• 区域1：[1,2,3,4,5]
• 区域2：[10,11,12,13]
• 区域3：[20, 21, 22, 23, 24, 25, 26, 27, 28]

处理规则 #

每个独立区域内的处理逻辑如下：

• 若b中某位置取值为1，将该位置后续的i个位置的值置为0
• 若位置原本取值为0，直接保持0值不变
• 完成i个位置的置零操作后不做额外处理，直到遇到下一个取值为1的位置，重复上述置零流程

结果示例 #

不同i取值对应的处理结果参考：

# i = 1
a     = [1, 2, 3, 4, 5, 10, 11, 12, 13, 20, 21, 22, 23, 24, 25, 26, 27, 28]
b_new = [0, 1, 0, 1, 0,  1,  0,  1,  0,  1,  0,  1,  0,  1,  0,  0,  0,  1]


# i = 2
a     = [1, 2, 3, 4, 5, 10, 11, 12, 13, 20, 21, 22, 23, 24, 25, 26, 27, 28]
b_new = [0, 1, 0, 0, 1,  1,  0,  0,  0,  1,  0,  0,  1,  0,  0,  0,  0,  1]


# i = 4
a     = [1, 2, 3, 4, 5, 10, 11, 12, 13, 20, 21, 22, 23, 24, 25, 26, 27, 28]
b_new = [0, 1, 0, 0, 0,  1,  0,  0,  0,  1,  0,  0,  0,  0,  0,  0,  0,  1]

现有前置实现 #

目前已经通过以下TensorFlow代码完成了a的区域分段划分，不确定该段代码是否具备复用价值：

a_shifted = tf.roll(a - 1, shift=-1, axis=0)
a_shifted_segs = tf.math.cumsum(tf.cast(a_shifted != a, dtype=tf.int64), exclusive=True)

# a_shifted_segs 输出结果：
# [0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2]

需求是找到能高效实现上述张量处理逻辑的方案。

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