We failed to lift variable creations out of this tf.function, so this tf.function cannot be run on XLA. A possible workaround is to move variable creation outside of the XLA compiled function.
最早用的pip -U 安装的keras没注意版本，直接可用。
之后装了一个第三方的Focal Loss库，结果自动把tf降了版本，后来再装keras只是==3.0结果就是这个版本不够新，导致了多卡分布式训练报xla错。折腾一下午，恍惚记得最早是3.0.5的keras，随后pip install keras==3.0.5，恢复正常。
