分析 JAX 代码：什么是 redzone_checker，为什么要花这么多时间？

Question

我找到了this post，但我仍然不清楚redzone_checker 内核在做什么以及为什么。具体来说，它是否应该占用我应用程序运行时间的 90% 以上？ TensorBoard 报告说它占用了我的 JAX 代码的绝大部分运行时间，我想知道

1. 实际上是这个内核占用了太多时间，还是这是使用 TensorBoard 分析 JAX 的副作用（即，输出在某种程度上具有误导性）？
2. 有没有办法减少redzone_checker 内核占用的时间？这是个好主意吗？

提前感谢您提供任何见解。

Answer 1

在分析之前确保预热。

可能是jit编译时间。