JAX，纯函数以及如何调试它

Question

我在调试 JAX 中非纯粹、非平凡的函数方面遇到了一些困难。他们的文档相当简单，当你开始学习它时就很好了。从概念上讲，纯函数相当容易理解。然而，当你开始开发一些新工具时，你可能会得到

Leaked tracers

并且调试起来非常困难。

如果能有一个有关 JAX 中重要示例的教程

Leaked tracer errors

以及这些示例的解决方案及其解释，那就太好了。

Answer 1

冒着误解您的问题的风险，因为它是相当开放式的，分别用于打印语句和交互式调试的

jax.debug.print

和

jax.debug.breakpoint

函数应该在纯和非纯 jax 函数中工作。

如果您分享代码片段，我们可以确定如何调试它！