使用TensorFlow ODE求解器获取梯度时出现ZeroDivisionError如何解决？

错误根因 #

从堆栈信息可以直接定位错误触发点是tf.concat算子的梯度计算逻辑，代码行concat_dim._numpy().item(0) % input_values[0]._rank()中，参与模运算的分母是拼接输入张量的秩，值为0时就会触发ZeroDivisionError，具体原因分两类：

• 你实现时变变量的计算逻辑里，存在对0秩标量张量做tf.concat的操作。把时变变量替换为常量时，该常量属于静态值，不会触发反向梯度路径里的concat梯度计算，因此不会报错；保留时变属性时，梯度反向传播会调用concat的梯度算子，就暴露了标量拼接的问题。
• TFP的ODE求解器默认使用伴随法（adjoint method）计算梯度，构造增广状态的过程中，时变变量的计算逻辑会产生0秩中间张量，触发拼接错误。

修复方案 #

• 优先排查时变变量计算逻辑里所有的tf.concat操作，给每个待拼接的张量增加秩校验tf.debugging.assert_rank_at_least(tensor, 1)，定位0秩张量的位置，对确实需要拼接的标量，先用tf.expand_dims(scalar, 0)扩展为1秩张量后再执行拼接操作。
• 你当前使用的TensorFlow 2.5.0、tensorflow-probability 0.13.0属于较老版本，该tf.concat梯度模零问题属于已知旧版本bug，升级到TensorFlow 2.8+、tensorflow-probability 0.16+版本可直接修复该问题。
• 暂时不想升级版本的话，可以修改ODE求解器的梯度计算方式，将gradient_algorithm参数设置为'forward'使用前向自动微分，规避伴随法梯度计算里的拼接逻辑，缺点是梯度计算速度会有明显下降。

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