我已经使用 Tensorflow 构建了一个 Monte-Carlo 模拟模型，该模型tf.while_loop用于执行迭代，我对使用tf.gradients. 该模型采用四个输入参数，因此tf.gradients返回四个偏导数（每个参数一个）。我正在尝试对tf.gradients计算衍生品的性能进行基准测试，但我的速度明显放缓。原始函数运行tf.gradients大约需要 0.03 秒，而计算导数大约需要 0.45 秒（15 倍减速）。我的期望是使用自动微分的梯度计算应该快得多。tf.gradients底层图表何时有问题tf.while_loop？
我正在使用带有 XLA 的 Tensorflow 2.5。