python-3.x - tanh 函数返回 'nan' 值；tanh 不稳定问题

Question

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

输出（具有意外值）

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. 我已经标记了占此输出的行（请参阅上面的代码）。

有什么解决办法吗？