如何转置不规则张量

Question

一切都在标题中。

我有一个 shape 的 RaggedTensor (nsamples, None, M1)，我想转置最后两个轴以获得 shape 的 RaggedTensor (nsamples, M1, None)。

tf.transpose不适用于参差不齐的张量，也不适用于tf.linalg.matrix_transpose（因为它们是相同的函数）。

有没有办法纯粹在 Tensorflow 中完成它，或者我必须for手动执行（慢速）循环？

Answer 1

首先，我希望这个答案能够在未来帮助一些人（包括我自己），尽管最初的问题已经晚了好几年了。

交换参差不齐张量的轴 1 和 2ragged_rank=1

对于手头的问题，可以使用 实现所需的转置tf.map_fn()。

def swap_dims_1_2(rt: tf.RaggedTensor):
    dtype = rt.flat_values.dtype
    
    def _fn(_t):
        return tf.RaggedTensor.from_tensor(tf.transpose(_t, [1, 0]))
    
    output_shape = (rt.flat_values.shape[1], None)
    return tf.map_fn(_fn, rt, fn_output_signature=tf.RaggedTensorSpec(shape=output_shape, dtype=dtype))

这是一个小测试：

data = [
    [[1, 2]], 
    [[3, 4], [5, 6], [7, 8], [9, 10]], 
    [[11, 12], [13, 14]],
]

rt: tf.RaggedTensor = tf.ragged.constant(data, ragged_rank=1)
print(rt.shape)  # (3, None, 2)
print(rt.nested_row_lengths())

r2 = swap_dims_1_2(rt)
print(r2)
print(r2.shape)  # (3, 2, None)
print(r2.nested_row_lengths())

我们还可以将不规则张量填充为常规张量，然后转置，这可能是不可取的，但可以用作检查上述函数的正确性：

# ... 
# `rt` was created same as above 

t = rt.to_tensor(default_value=0)  # shape (3, 4, 2)
t_transpose = tf.transpose(t, perm=[0, 2, 1])  # shape (3, 2, 4)

rt_transpose = tf.ragged.boolean_mask(t_transpose, mask=t_transpose>0)
rt_transpose._set_shape([None, 2, None])
print(rt_transpose)
print(rt_transpose.shape)  # (3, 2, None)
print(rt_transpose.nested_row_lengths())

关于转置不规则张量的注意事项

一般来说，在不填充一些填充值的情况下，不可能在不规则张量中转置维度，特别是当涉及多个不规则维度时。我们可以这样想：张量，无论是否参差不齐，都是列表的嵌套列表。转置张量意味着重新排序嵌套级别，并且生成的张量的内部列表可以以统一的步长包含原始张量的元素。对于参差不齐的张量来说，这会成为问题，因为对于某些维度，“统一步长的元素”没有明确定义。

因此，当考虑“转置不规则张量”或“交换不规则张量中的轴”时，最好写下一个例子，看看假设的转置操作是否真的达到了预期的效果。