[yoch]

muzero_policy(..., dirichlet_fraction=0.0) still samples the Dirichlet noise (gamma sampling, a few while loops in the compiled helper) and multiplies it by zero. The same happens in gumbel_muzero_policy with gumbel_scale=0.0, which is the documented setting for evaluating perfect-information games.

When the coefficient is a Python zero (static), the helper now returns the probs, or zero Gumbels, directly. The random key is still split in the same way, so the other keys are unchanged, and a traced zero keeps the old path. On a CPU with batch 128 and 81 actions, a whole 8-simulation muzero_policy call with dirichlet_fraction=0.0 goes from 22.5 ms to 12.3 ms (best of 7); with a non-zero fraction the time is unchanged (22.1 ms against 21.5 ms), as expected. The gain is only visible on short searches, where the sampling is a large share of the call.

One detail worth a look in review: the old mixing (1 - fraction) * probs + fraction * noise promotes bfloat16 or float16 probs to the default float type. The shortcut keeps that promotion (it casts the probs without sampling), so outputs and dtypes match a traced zero fraction. The tests check this for bfloat16, float16 and float32 priors, on the helper and on the full tree of muzero_policy.

Tests: pytest from _testing/: 29 passed, pylint clean on the changed source file (the one remaining line-too-long message in policies_test.py is a line that is already on main).

On a Colab T4 the gain is small (a whole 8-simulation call at batch 128 and 81 actions: 5.9 ms against 6.0 ms with dirichlet_fraction=0.0, medians of 5; 5.4 ms against 6.0 ms for the best call) and there is no change with a non-zero fraction (5.8 ms against 6.3 ms, same trees). On a TPU v5e the same call goes from 3.39 ms to 3.24 ms (medians of 5; 3.41 against 3.26 with a non-zero fraction), so the benefit is mostly on CPU and for short searches; I see no cost elsewhere.

Possible overlap: #123 and #121 also edit policies.py; both apply cleanly on top of this branch.