tanh 函数返回“nan”值； tanh 不稳定问题

如何解决tanh 函数返回“nan”值； tanh 不稳定问题

这是一个简单的代码，用于测试来自 Deepmind 的新 TanH 包的 distrax 功能。

输出（带有意外值）

WARNING:absl:No GPU/TPU found,falling back to cpu. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

[nan nan nan nan nan nan nan nan]
[nan nan nan nan nan nan nan nan]

代码

import distrax
import jax
import jax.numpy as jnp

BATCH_SIZE = 8
means = jnp.zeros((BATCH_SIZE,))
log_stds = jnp.full_like(means,10)

base_dist = distrax.normal(
    loc=means,scale=jnp.exp(log_stds))

dist = distrax.Transformed(distribution=base_dist,bijector=distrax.Tanh()) ----> problem with this line

@jax.jit
def test(seed):
  samples = dist.sample(seed=seed)
  return dist.log_prob(samples)

print(test(seed=jax.random.PRNGKey(1)))
print(test(seed=jax.random.PRNGKey(2)))

其他上下文

1. 我在 jupyter notebook（python3 内核）中运行了这个程序
2. 我已经标记了说明此输出的行（请参阅上面的代码）。

有什么解决办法吗？