jax-ml/jax · error · ValueError

dtype argument to `gumbel` must be a float dtype, got {dtype

Error message

dtype argument to `gumbel` must be a float dtype, got {dtype}

What it means

jax.random.gumbel requires a floating-point output dtype because Gumbel samples are produced with log/exp float math. Integer or complex dtypes raise ValueError before tracing.

Source

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

      with mode="high" this is increased to ~32, at approximately double the
      computational cost.
    out_sharding: Optional. Specifies how the output array should be sharded
      across devices in multi-device computation. Can be a
      :class:`~jax.sharding.NamedSharding`, a :class:`~jax.sharding.PartitionSpec`
      (``P``), or ``None`` (default). When specified, the output will be sharded
      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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass jnp.float32/jnp.float64 or omit dtype (defaults to float).
2. Keep the Gumbel noise float and cast only the resulting argmax indices to int afterwards.
3. Validate configurable dtypes against np.floating before the call.

Example fix

// before
g = jax.random.gumbel(key, (10, 4), dtype=jnp.int32)

// after
g = jax.random.gumbel(key, (10, 4), dtype=jnp.float32)

from jax._src import dtypes
assert dtypes.issubdtype(dtypes.check_and_canonicalize_user_dtype(dtype or float), np.floating)

def is_float_dtype(dtype) -> bool:
    from jax._src import dtypes
    import numpy as np
    return dtypes.issubdtype(dtypes.check_and_canonicalize_user_dtype(dtype or float), np.floating)

When it happens

Trigger: jax.random.gumbel(key, shape, dtype=jnp.int32) or any dtype where dtypes.issubdtype(dtype, np.floating) is False.

Common situations: Using gumbel for Gumbel-max trick over discrete choices and mistakenly assuming the output dtype should match the discrete labels; sharing a dtype config across samplers.

Related errors

• dtype argument to `exponential` must be a float dtype, got {
• dtype argument to `gamma` must be a float dtype, got {dtype}
• Must provide valid mode for gumbel got: %s
• dtype argument to `laplace` must be a float dtype, got {dtyp
• dtype argument to `logistic` must be a float dtype, got {dty

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