使用tf.experimental.ExtensionType时TensorFlow追踪异常问题咨询

核心原因 #

tf.function是否触发重追踪，完全取决于输入的**追踪类型（Tracing Type）**是否发生变化。你使用tf.experimental.ExtensionType时出现不符合预期的情况，大概率是以下几个原因：

1. ExtensionType字段未正确标记为可追踪类型 #

如果你的Apple类中flavor字段的类型标注不是tf.Tensor（比如写成list、未标注类型，或者用了其他非TensorFlow可追踪类型），TensorFlow会把它当作普通Python对象处理，不会监测其内部张量的形状、dtype变化。这种情况下，哪怕你修改了flavor的张量属性，输入的追踪类型哈希值不会变，自然不会触发重追踪。

2. 错误重写了__tf_tracing_type__方法 #

如果为了实现单次追踪手动重写了__tf_tracing_type__，但覆盖了ExtensionType的默认逻辑，导致flavor的张量元信息（形状、dtype）没有被纳入追踪类型的计算中。这样不管flavor怎么变，追踪类型都保持一致，tf.function就只会追踪一次。

3. 未创建新的ExtensionType实例 #

tf.experimental.ExtensionType是不可变类型，实例创建后无法修改其属性值。如果你尝试直接修改现有Apple实例的flavor，实际上不会生效，传入tf.function的还是原来的实例，自然不会触发重追踪。

解决方法 #

针对上述问题，你可以按以下步骤调整：

• 正确标注字段类型：确保Apple类的flavor字段明确标注为tf.Tensor，让TensorFlow自动追踪其元信息：

  class Apple(tf.experimental.ExtensionType):
      flavor: tf.Tensor
  

• 保留默认追踪逻辑：不要手动重写__tf_tracing_type__，ExtensionType的默认实现会自动将所有可追踪字段的追踪类型（包括张量的形状、dtype）组合成自身的追踪类型。
• 创建新实例传入：修改flavor的形状或dtype后，必须创建新的Apple实例传入tf.function，比如：

  # 第一次调用
  apple1 = Apple(tf.constant([1, 2]))
  get_mixed_flavor(apple1)
  
  # 修改形状后创建新实例
  apple2 = Apple(tf.constant([[1, 2]]))
  get_mixed_flavor(apple2)  # 此时会触发重追踪
  

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