TensorFlow AutoGraph多输出多项式模型训练张量维度报错排查

• 维度越界的核心触发源是*tf.vectorized_map（vmap）对嵌套多输出结构的切片逻辑错位*：AutoGraph编译训练流程时，会自动插入vmap节点实现批量向量化。单输出场景下，vmap会将输出张量的第一维识别为批量维并剥离，下游函数接收到的张量维度和单样本推理完全对齐；当模型输出为元组/字典类嵌套结构时，vmap会将输出结构的嵌套层级误纳入切片偏移量计算，最终传入多项式计算函数的系数张量被错误压平为1维向量——原本预期的单样本2D系数矩阵（[多项式数量, 系数阶数]）丢失了结构维度，此时硬编码第二维索引的Horner逻辑自然触发StridedShape不匹配的越界报错。
• 手动推理、单样本手动计算无报错的原因：该流程未经过AutoGraph的自动vmap插入步骤，张量维度完全按照手动调用逻辑构建，不存在切片偏移问题。
• 移除第二组输出即可正常训练的现象，完全佐证了上述机制：单张量输出不会触发vmap的嵌套结构错位，批量维的剥离过程和单样本维度完全对齐，因此不会出现偏移。

• 修正Horner函数的索引逻辑，适配动态维度
  若原函数采用coeff[:, order]这类硬绑定第二维的索引方式，在张量被压平时必然触发越界。将索引替换为兼容所有前置维度的省略号语法，保证无论前置批量维数量如何变化，都能正确沿系数阶数维度取值，同时完全保留AutoGrad求导能力：

  def horner_polyval(coeff, x):
      # 兼容任意前置批量维，coeff最后一维为系数阶数
      res = coeff[..., 0]
      for order_idx in range(1, coeff.shape[-1]):
          res = res * x + coeff[..., order_idx]
      return res
  

• 调整模型输出结构，消除vmap偏移触发条件
  将原本输出的两个嵌套分支（系数矩阵、输入副本）沿最后一维拼接为单张量输出，在损失函数内部按固定维度位置拆分取值。单张量输出不会触发vmap的嵌套切片错位逻辑，批量维的剥离过程和单样本推理完全对齐，从根源上避免维度偏移。
• 自定义训练循环，规避自动vmap插入
  若需保留原有嵌套多输出结构，可弃用model.fit内置的自动训练流程，改为手动实现带梯度计算的训练步，不经过AutoGraph自动插入的vmap节点，完全手动控制张量维度流向：

  @tf.function
  def train_step(batch_x, batch_label):
      with tf.GradientTape() as tape:
          # 直接调用模型获取嵌套输出，手动拆分计算
          coeff_mat, input_copy = model(batch_x, training=True)
          pred = horner_polyval(coeff_mat, input_copy)
          loss = tf.reduce_mean(tf.square(pred - batch_label))
      grads = tape.gradient(loss, model.trainable_weights)
      optimizer.apply_gradients(zip(grads, model.trainable_weights))
      return loss
  

• 快速定位维度错位的调试方法
  在多项式计算函数入口加入形状打印逻辑，对比单样本推理和训练流程下传入的系数张量形状：如果训练流程下的形状比预期少1个维度，即可直接确认是vmap的嵌套切片偏移导致的维度错位，快速定位偏移位置。

内容的提问来源于stack exchange，提问作者Alex R.