JAX如何处理返回为nan

Zhang庆欢

于 2022-08-17 18:48:38 发布

阅读量395

点赞数

CC 4.0 BY-SA版权

分类专栏： Jax pythonBug 文章标签： numpy python

8 篇文章

订阅专栏

4 篇文章

订阅专栏

当在JAX中遇到除以零导致的nan值时，定位问题可能变得困难。解决方法是在代码开始前加入`from jax.config import config; config.update(jax_debug_nans True)`，这将帮助你在出错的位置得到更具体的错误信息，而不是简单的nan返回。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

1.问题描述

使用JAX时返回了nan：

from jax import numpy as jnp
x = jnp.divide(0., 0.)
print("x value:", x)

显示如下：
在这里插入图片描述
这个时候其实没办法定位问题在哪里，因为你写的代码是成百上千行的！！！

from jax import numpy as jnp
from jax.config import config
config.update("jax_debug_nans", True)

x = jnp.divide(0., 0.)
print("x value:", x)

显示如下：
在这里插入图片描述

因此可以加上上面的语法以定位发生错误的位置，而不是返回一个nan值！