处理动态数量动态形状张量列表的tf.function签名配置方法

RaggedTensor 方案报错本质 #

你遇到的Spec不兼容问题，核心是对RaggedTensorSpec的不规则轴推导逻辑理解偏差：

• 报错提示的ragged_rank数值不匹配：你用RaggedTensor.from_row_lengths构造的张量ragged_rank=1，对应形状[2, None, 3]——第0轴是固定长度的列表维度，第1轴是不规则的行维度；但你在input_signature里写shape=[None, None, 3]时没有显式指定ragged_rank，TF会自动按形状里的None数量推导为ragged_rank=2，默认两个None轴都是不规则轴，和实际传入的张量结构完全不匹配。
• 此前认为“RaggedTensorSpec的None无法标记非不规则的动态维度”是认知偏差：只要显式指定ragged_rank参数，就能清晰区分“普通动态固定维度（仅维度值运行时确定，轴上元素无形状差异）”和“不规则维度（轴上元素形状可变）”。
  适配你需求的正确Spec写法为：

tf.RaggedTensorSpec(shape=[None, None, 3], dtype=tf.float32, ragged_rank=1)

上述配置中，第一个None对应动态的列表长度（非不规则轴，仅长度运行时确定），第二个None对应不规则的行维度（每个张量的行数运行时可变），完全匹配你的输入场景。

TensorArray 方案报错本质 #

你遇到的TypeError是跨函数传递TensorArray对象的兼容性问题，并非TensorArraySpec完全无法用于该场景：

• 旧版TF的tf.function入参校验逻辑，不支持将裸TensorArray对象作为单参数传入，会错误匹配input_signature的序列结构，触发“入参必须为序列”的校验错误。
• 即使在新版本TF中，直接跨函数传递TensorArray也容易因为构造参数（如dynamic_size、infer_shape）和Spec默认值不匹配触发校验失败，没必要在图执行时跨边界传递TensorArray对象。

无重复追踪的可落地实现 #

两种实现路径均可满足需求，且不会因为输入的列表长度、元素行维度变化触发重复追踪：

方案1：修正RaggedTensor的Spec配置 #

import tensorflow as tf

a = tf.constant([[1, 2, 3], [4, 5, 6]], dtype=tf.float32)
b = tf.constant([[7, 8, 9]], dtype=tf.float32)

values = tf.concat([a, b], axis=0)
lens = tf.stack([tf.shape(t)[0] for t in [a, b]])
ragged_tensor = tf.RaggedTensor.from_row_lengths(values, lens)

@tf.function(input_signature=(
    tf.RaggedTensorSpec(shape=[None, None, 3], dtype=tf.float32, ragged_rank=1),
))
def myFunc(x):
    return tf.math.reduce_sum(x)

# 正常执行，所有符合“列数固定为3、dtype为float32”的输入复用同一计算图
print(myFunc(ragged_tensor))

方案2：函数内部构造TensorArray的兼容写法 #

把TensorArray的构造逻辑移到tf.function内部，仅传基础张量作为入参，完全规避跨对象传递的校验问题：

@tf.function(input_signature=(
    # 所有元素拼接后的扁平张量，最后一维固定为3
    tf.TensorSpec(shape=[None, 3], dtype=tf.float32),
    # 每个元素的行长度标记，长度动态
    tf.TensorSpec(shape=[None], dtype=tf.int32)
))
def myFunc1(flat_values, row_lengths):
    # 函数内部动态构造TensorArray
    array = tf.TensorArray(
        tf.float32,
        size=tf.shape(row_lengths)[0],
        dynamic_size=True,
        infer_shape=False,
        element_shape=[None, 3]
    )
    # 按长度拆分写入元素
    start = 0
    for i in tf.range(tf.shape(row_lengths)[0]):
        cur_len = row_lengths[i]
        array = array.write(i, flat_values[start:start+cur_len])
        start += cur_len
    return tf.math.reduce_sum(array.concat())

# 正常执行
print(myFunc1(values, lens))

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