对每列不同行数的张量求和。

Question

我希望能够在Tensorflow中对一个张量进行还原求和，对于每一列，我们只对行的一个子集进行求和。为了说明这一点，请考虑以下内容

import tensorflow as tf


X = tf.constant(
    [
        [1, 3, 2],
        [0, 5, 8],
        [1, 6, 2]
    ],
    tf.float32
)

row_max = tf.constant([3, 2, 1], tf.int64)

然后，我想在Tensorflow中做以下操作，使梯度可以流动。

partial_sum = 0.0

for col_idx in range(X.shape[1]):
    partial_sum += tf.reduce_sum(X[:row_max[col_idx], col_idx]])

这样我就能得到 1+0+1+3+5+2 = 12

然而，我不知道如何在Tensorflow中做到这一点。我研究了很多不同的方法。tf.ragged.range, tf.segment_sum 等。我认为 tf.gather_nd 可以工作，但即使如此，我也不知道如何建立索引张量。在Numpy中，我可以做这样的事情。

import numpy as np


X_np = X.numpy()

idx0 = np.concatenate(
    [
        i * np.ones(row_max[i])
        for i in range X_np.shape[1]
    ],
    axis=0
).astype(np.int64)

idx1 = np.concatenate(
    [
        np.arange(row_max[i])
        for i in range X_np.shape[1]
    ],
    axis=0
).astype(np.int64)


X_np[idx0, idx1].sum()

在Tensorflow中实现我的目标的最好方法是什么？

Answer 1

下面是一个简单的方法。

import tensorflow as tf

x = tf.constant(
    [
        [1, 3, 2],
        [0, 5, 8],
        [1, 6, 2]
    ],
    tf.float32
)
row_max = tf.constant([3, 2, 1], tf.int64)

# Make mask for each column
row_idx = tf.range(tf.shape(x, out_type=row_max.dtype)[0])
mask = tf.expand_dims(row_idx, 1) < row_max
mask_f = tf.dtypes.cast(mask, x.dtype)
# Mask elements and sum
result = tf.reduce_sum(mask_f * x)
tf.print(result)
# 12

# Alternatively, you can mask the elements and sum
result = tf.reduce_sum(tf.boolean_mask(x, mask))