jax-ml/jax · error · ValueError

Must provide valid mode for gumbel got: %s

Error message

Must provide valid mode for gumbel got: %s

What it means

jax.random.gumbel's mode argument controls the dynamic range of the sampler and must be one of 'highest', 'high', or 'low'. When mode=None it defaults to 'high' or 'low' based on the jax_use_high_dynamic_range_gumbel config, but any explicitly supplied value outside the allowed trio raises ValueError.

Source

Thrown at jax/_src/random/core.py:2277

      according to the given sharding specification. Primarily used in explicit
      sharding mode.
      See the `explicit sharding tutorial <https://docs.jax.dev/en/latest/parallel.html>`_
      for more details.

  Returns:
    A random array with the specified shape and dtype.
  """
  key, _ = _check_prng_key("gumbel", key)
  dtype = dtypes.check_and_canonicalize_user_dtype(
      float if dtype is None else dtype)
  if not dtypes.issubdtype(dtype, np.floating):
    raise ValueError(f"dtype argument to `gumbel` must be a float "
                     f"dtype, got {dtype}")
  shape = core.canonicalize_shape(shape)
  if mode is None:
    mode = "high" if config.use_high_dynamic_range_gumbel.value else "low"
  if mode not in ("highest", "high", "low"):
    raise ValueError("Must provide valid mode for gumbel got: %s" % mode)
  out_sharding = canonicalize_sharding_for_samplers(out_sharding, "gumbel", shape)
  return maybe_auto_axes(_gumbel, out_sharding, shape=shape, dtype=dtype,
                         mode=mode)(key)

def _safe_int_to_float(bits, dtype):
  """Converts bits: u32[2,...] into f32[...] in the range (0,1)."""
  if bits.dtype != np.uint32 or dtype != np.float32:
    raise RuntimeError("_safe_int_to_float only works for u32 -> f32")
  finfo = dtypes.finfo(dtype)
  hiclz, loclz = lax.clz(bits)
  hi, lo = bits

  mantissa = lax.bitwise_or(
      lax.shift_left(hi, hiclz),
      jnp.where(
          hiclz == 32,
          lax.shift_left(lo, loclz),
          lax.shift_right_logical(lo, finfo.bits - hiclz)))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use 'high', 'low', or 'highest' exactly (lowercase).
2. Omit mode to get the configuration-driven default ('high' if use_high_dynamic_range_gumbel is enabled, else 'low').
3. Validate mode against ('highest','high','low') in config-loading code.

Example fix

// before
g = jax.random.gumbel(key, (8,), mode='max')

// after
g = jax.random.gumbel(key, (8,), mode='high')

if mode is not None:
    assert mode in ('highest', 'high', 'low'), f'invalid gumbel mode: {mode}'

def is_valid_gumbel_mode(mode) -> bool:
    return mode is None or mode in ('highest', 'high', 'low')

When it happens

Trigger: jax.random.gumbel(key, shape, mode='max'), mode='HIGH' (case mismatch), or passing a non-string; only when mode is explicitly given (None is replaced by the config default).

Common situations: Newer mode argument unfamiliar to users porting old code; typos; passing the mode through from a config that never validated it.

Related errors

• dtype argument to `gumbel` must be a float dtype, got {dtype
• dirichlet requires alpha.ndim >= 1, got alpha.ndim == {}
• dtype argument to `exponential` must be a float dtype, got {
• method argument to `gamma` must be one of {'exact', 'approxi
• dtype argument to `gamma` must be a float dtype, got {dtype}

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/bb9b9791c468f77e. Report an issue: GitHub.