def test_mix_value_with_zero_visits(self):
    """Tests that zero visit counts do not divide by zero."""
    raw_value = jnp.array(-0.8)
    prior_logits = jnp.array([-jnp.inf, -1.0, 2.0, -jnp.inf])
    probs = jax.nn.softmax(prior_logits)
    visit_counts = jnp.array([0, 0, 0, 0])
    qvalues = jnp.zeros_like(probs)
    with jax.debug_nans():
      mix_value = qtransforms._compute_mixed_value(
          raw_value, qvalues, visit_counts, probs)

    np.testing.assert_allclose(raw_value, mix_value)
Tests that zero visit counts do not divide by zero.
mctx/_src/tests/qtransforms_test.py