{"record":{"id":"b6eeeb791e5e37fe","repo":"jax-ml/jax","slug":"dtype-argument-to-gumbel-must-be-a-float-dtype","errorCode":null,"errorMessage":"dtype argument to `gumbel` must be a float dtype, got {dtype}","messagePattern":"dtype argument to `gumbel` must be a float dtype, got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/random/core.py","lineNumber":2271,"sourceCode":"      with mode=\"high\" this is increased to ~32, at approximately double the\n      computational cost.\n    out_sharding: Optional. Specifies how the output array should be sharded\n      across devices in multi-device computation. Can be a\n      :class:`~jax.sharding.NamedSharding`, a :class:`~jax.sharding.PartitionSpec`\n      (``P``), or ``None`` (default). When specified, the output will be sharded\n      according to the given sharding specification. Primarily used in explicit\n      sharding mode.\n      See the `explicit sharding tutorial <https://docs.jax.dev/en/latest/parallel.html>`_\n      for more details.\n\n  Returns:\n    A random array with the specified shape and dtype.\n  \"\"\"\n  key, _ = _check_prng_key(\"gumbel\", key)\n  dtype = dtypes.check_and_canonicalize_user_dtype(\n      float if dtype is None else dtype)\n  if not dtypes.issubdtype(dtype, np.floating):\n    raise ValueError(f\"dtype argument to `gumbel` must be a float \"\n                     f\"dtype, got {dtype}\")\n  shape = core.canonicalize_shape(shape)\n  if mode is None:\n    mode = \"high\" if config.use_high_dynamic_range_gumbel.value else \"low\"\n  if mode not in (\"highest\", \"high\", \"low\"):\n    raise ValueError(\"Must provide valid mode for gumbel got: %s\" % mode)\n  out_sharding = canonicalize_sharding_for_samplers(out_sharding, \"gumbel\", shape)\n  return maybe_auto_axes(_gumbel, out_sharding, shape=shape, dtype=dtype,\n                         mode=mode)(key)\n\ndef _safe_int_to_float(bits, dtype):\n  \"\"\"Converts bits: u32[2,...] into f32[...] in the range (0,1).\"\"\"\n  if bits.dtype != np.uint32 or dtype != np.float32:\n    raise RuntimeError(\"_safe_int_to_float only works for u32 -> f32\")\n  finfo = dtypes.finfo(dtype)\n  hiclz, loclz = lax.clz(bits)\n  hi, lo = bits\n","sourceCodeStart":2253,"sourceCodeEnd":2289,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/random/core.py#L2253-L2289","documentation":"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.","triggerScenarios":"jax.random.gumbel(key, shape, dtype=jnp.int32) or any dtype where dtypes.issubdtype(dtype, np.floating) is False.","commonSituations":"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.","solutions":["Pass jnp.float32/jnp.float64 or omit dtype (defaults to float).","Keep the Gumbel noise float and cast only the resulting argmax indices to int afterwards.","Validate configurable dtypes against np.floating before the call."],"exampleFix":"// before\ng = jax.random.gumbel(key, (10, 4), dtype=jnp.int32)\n\n// after\ng = jax.random.gumbel(key, (10, 4), dtype=jnp.float32)","handlingStrategy":"type-guard","validationCode":"from jax._src import dtypes\nassert dtypes.issubdtype(dtypes.check_and_canonicalize_user_dtype(dtype or float), np.floating)","typeGuard":"def is_float_dtype(dtype) -> bool:\n    from jax._src import dtypes\n    import numpy as np\n    return dtypes.issubdtype(dtypes.check_and_canonicalize_user_dtype(dtype or float), np.floating)","tryCatchPattern":null,"preventionTips":["Keep Gumbel noise float; cast argmax indices to int instead."],"tags":["jax","random","gumbel","dtype","input-validation"],"backgroundTag":"invalid-dtype-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}