[jit] jax_high_dynamic_range_gumbel is missing from the JIT cache key

Description

Changing jax_high_dynamic_range_gumbel at runtime does not affect an already compiled jax.jit callable that invokes jax.random.gumbel. The callable silently continues using the sampling mode active during its first compilation.

This is reproducible in both directions (low -> high and high -> low) and on CPU and GPU.

Minimal reproducer

import numpy as np
import jax
import jax.random as jr

FLAG = "jax_high_dynamic_range_gumbel"
KEY = jr.PRNGKey(0)

jax.config.update(FLAG, False)
jax.clear_caches()

f = jax.jit(lambda key: jr.gumbel(key, (4,)))

low = np.asarray(f(KEY))

# Change the documented runtime configuration without clearing caches.
jax.config.update(FLAG, True)
after_update = np.asarray(f(KEY))

print("low:                       ", low)
print("same callable, flag=True:  ", after_update)
print("cache size:                ", f._cache_size())

# Force a fresh compilation under flag=True.
jax.clear_caches()
fresh_high = np.asarray(f(KEY))

print("after jax.clear_caches():  ", fresh_high)

# These assertions demonstrate the current behavior.
assert np.array_equal(after_update, low)
assert not np.allclose(fresh_high, low)

Observed output:

low:                        [ 2.9233725   3.8326178  -0.09689324  0.27725708]
same callable, flag=True:   [ 2.9233725   3.8326178  -0.09689324  0.27725708]
cache size:                 1
after jax.clear_caches():   [-1.0818486  -1.3463638   0.906579    0.45828176]

The reverse direction behaves the same way: a callable first compiled with the flag set to True continues producing high-mode samples after the flag is changed to False, until the cache is cleared.

Expected behavior

Changing jax_high_dynamic_range_gumbel should invalidate compiled functions whose traced computation depends on the flag. The next invocation should use the newly selected Gumbel sampling mode.

Suspected root cause

jax.random.gumbel reads the configuration during tracing and converts it into a static mode argument:

if mode is None:
  mode = "high" if config.use_high_dynamic_range_gumbel.value else "low"

return maybe_auto_axes(
    _gumbel, out_sharding, shape=shape, dtype=dtype, mode=mode
)(key)

However, the configuration state has include_in_trace_context=True but not include_in_jit_key=True:

use_high_dynamic_range_gumbel = bool_state(
    name="jax_high_dynamic_range_gumbel",
    default=False,
    help="...",
    include_in_trace_context=True,
)

Consequently, the two flag values produce different jaxprs, but the outer JIT cache can reuse an executable traced under the previous value.

Adding include_in_jit_key=True locally causes the configuration change to retrace the callable: its cache size changes from 1 to 2 and the output immediately switches to the requested sampling mode.

This is consistent with the treatment of other trace-sensitive configuration states, such as jax_softmax_custom_jvp and jax_threefry_gpu_kernel_lowering.

Versions tested

The reproducer was confirmed with:

• JAX/JAXLIB 0.10.2 on CPU and NVIDIA GPU
• JAX/JAXLIB 0.11.1 on CPU
• JAX main at bd135887c682c92aa121216ed17d076e1d5d5be0 (0.11.2.dev20260910) with JAXLIB 0.11.1 on CPU

No existing issue describing this cache-key bug was found. PR #26870 originally introduced the high-dynamic-range Gumbel option, but does not address runtime JIT cache invalidation.

Workarounds

• Pass mode="high" or mode="low" explicitly to jax.random.gumbel.
• Set the configuration before compiling any relevant functions.
• Call jax.clear_caches() after changing the configuration.

System info (python version, jaxlib version, accelerator, etc.)

jax: 0.10.2 jaxlib: 0.10.2 numpy: 2.5.1 python: 3.12.13 | packaged by Anaconda, Inc. | (main, Jul 9 2026, 14:38:16) [GCC 14.3.0] device info: NVIDIA GeForce RTX 4090-1, 1 local devices" process_count: 1 platform: uname_result(system='Linux', node='DESKTOP-09NMRPF', release='6.6.87.2-microsoft-standard-WSL2', version='#1 SMP PREEMPT_DYNAMIC Thu Jun 5 18:30:46 UTC 2025', machine='x86_64') XLA_PYTHON_CLIENT_PREALLOCATE=false

$ nvidia-smi Fri Sep 11 01:55:19 2026 NVIDIA-SMI 590.57 Driver Version: 591.86 CUDA Version: 13.1 GPU 0: NVIDIA GeForce RTX 4090, 24564 MiB