def test_mask_all_invalid_actions(self):
    """Tests a state with no valid action."""
    logits = jnp.array([-jnp.inf, -jnp.inf, -jnp.inf, -jnp.inf])
    invalid_actions = jnp.array([1.0, 1.0, 1.0, 1.0])
    masked_logits = policies._mask_invalid_actions(
        logits, invalid_actions)
    np.testing.assert_allclose(
        jnp.array([0.25, 0.25, 0.25, 0.25]),
        jax.nn.softmax(masked_logits))
Tests a state with no valid action.
mctx/_src/tests/policies_test.py