使用Cohen's Kappa指标时TPU训练失败但CPU运行正常的问题求助

TPU依赖XLA编译器对计算图做静态优化，要求所有控制流分支（比如tf.cond）的输出形状必须完全一致，CPU/GPU对维度不匹配的张量支持隐式广播兼容，因此不会触发报错。
你遇到的报错是tfa实现的CohenKappa中_update_confusion_matrix方法内部的条件分支，两个分支的输出一个带[1, <=4]的二维形状，一个是[<=4]的一维形状，XLA编译阶段检测到形状不匹配直接抛出错误。
你观测到的动态输入日志确实是干扰项，只要输入的batch维度是动态的都会输出该提示，和本问题无关。

• 方案1：自定义TPU兼容的CohenKappa指标
  直接重载tfa.metrics.CohenKappa的_update_confusion_matrix和result方法，核心修改点是统一所有条件分支的输出形状，提前固定类别数避免运行时动态推导：

import tensorflow as tf
import tensorflow_addons as tfa

class TPUCompatibleCohenKappa(tfa.metrics.CohenKappa):
    def _update_confusion_matrix(self, y_true, y_pred, sample_weight):
        # 提前固定num_classes，不要动态推断
        num_classes = self.num_classes
        y_true = tf.cast(y_true, tf.int32)
        y_pred = tf.cast(y_pred, tf.int32)
        # 统一分支输出形状，所有输出都去掉多余的维度
        cm = tf.math.confusion_matrix(
            y_true,
            y_pred,
            num_classes=num_classes,
            weights=sample_weight,
            dtype=self.dtype
        )
        # 强制形状对齐，避免动态轴问题
        cm = tf.ensure_shape(cm, (num_classes, num_classes))
        self.confusion_mtx.assign_add(cm)
    
    def result(self):
        # 原result方法如果有条件分支也做同样的形状统一处理
        cm = self.confusion_mtx
        n = tf.reduce_sum(cm)
        sum0 = tf.reduce_sum(cm, axis=0)
        sum1 = tf.reduce_sum(cm, axis=1)
        expected = tf.reduce_sum(sum0 * sum1) / n
        w_mat = self.weight_matrix
        if self.weightage is not None:
            w_mat = self._get_weight_matrix(self.num_classes)
        k = tf.reduce_sum(w_mat * cm)
        expected_k = tf.reduce_sum(w_mat * sum1[:, tf.newaxis] * sum0[tf.newaxis, :]) / n
        return 1 - (k / expected_k)

使用时直接用TPUCompatibleCohenKappa替换原有tfa.metrics.CohenKappa即可。

• 方案2：离线计算CohenKappa
  TPU训练阶段仅使用loss、准确率等原生TPU兼容的指标，训练完成后导出模型权重，在CPU/GPU环境加载权重后离线计算CohenKappa指标，不需要修改训练阶段的计算逻辑。

• 在CPU环境开启XLA编译即可复现该报错，无需反复提交TPU任务调试：tf.config.optimizer.set_jit(True)
• 给所有tf.cond分支的输出添加tf.ensure_shape手动指定静态形状，强制所有分支输出完全一致
• 避免在条件分支内部做改变张量维度的操作（比如squeeze、expand_dims），统一在分支外部做维度调整

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