{"record":{"id":"4e9a5717e7326109","repo":"jax-ml/jax","slug":"dtype-argument-to-geometric-must-be-an-int-dtype","errorCode":null,"errorMessage":"dtype argument to `geometric` must be an int dtype, got {dtype}","messagePattern":"dtype argument to `geometric` must be an int dtype, got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/random/core.py","lineNumber":3328,"sourceCode":"      jax_enable_x64 is true, otherwise int32).\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 dtype and with shape given by ``shape`` if\n    ``shape`` is not None, or else by ``p.shape``.\n  \"\"\"\n  key, _ = _check_prng_key(\"geometric\", key)\n  dtype = dtypes.check_and_canonicalize_user_dtype(\n      int if dtype is None else dtype)\n  if not dtypes.issubdtype(dtype, np.integer):\n    raise ValueError(\"dtype argument to `geometric` must be an int \"\n                     f\"dtype, got {dtype}\")\n  shape = _check_broadcast_shapes(\"geometric\", shape, p)\n  out_sharding = canonicalize_sharding_for_samplers(out_sharding, \"geometric\", shape)\n  return _geometric(key, p, shape, dtype, out_sharding)\n\n@jit(static_argnums=(2, 3, 4))\ndef _geometric(key, p, shape, dtype, out_sharding) -> Array:\n  check_arraylike(\"geometric\", p)\n  p, = promote_dtypes_inexact(p)\n  u = uniform(key, shape, p.dtype, out_sharding=out_sharding)\n  # TODO(jakevdp): switch to log_u = lax.log1p(u - 1)\n  # For now we map u=0 to u=1 to avoid inf in log_u without otherwise\n  # changing samples produced for a given key.\n  u = jnp.where(u == 0, 1, u)\n  log_u = lax.log(u)\n  log_one_minus_p = lax.log1p(-p)\n  log_one_minus_p = jnp.broadcast_to(log_one_minus_p, shape, out_sharding=out_sharding)\n  g = lax.floor(lax.div(log_u, log_one_minus_p)) + 1","sourceCodeStart":3310,"sourceCodeEnd":3346,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/random/core.py#L3310-L3346","documentation":"jax.random.geometric is a discrete sampler, so it requires an integer dtype (default int). After canonicalization, dtypes.issubdtype(dtype, np.integer) is checked; passing float, bool, bfloat16, or complex dtypes raises this ValueError.","triggerScenarios":"Calling jax.random.geometric(key, p, shape, dtype=np.float32) or any non-integer dtype.","commonSituations":"Using a float dtype out of habit from continuous samplers; np.bool_ is not np.integer so boolean requests also fail.","solutions":["Pass an int dtype such as np.int32 or jnp.uint32, or omit dtype","If floats are needed downstream, sample as int then cast: jnp.asarray(x, np.float32)"],"exampleFix":"// before\ng = jax.random.geometric(key, 0.3, dtype=jnp.float32)\n// after\ng = jax.random.geometric(key, 0.3, dtype=jnp.int32)\n# if floats needed:\ngf = jax.random.geometric(key, 0.3, dtype=jnp.int32).astype(jnp.float32)","handlingStrategy":"validation","validationCode":"import numpy as np\nassert dtype is None or np.issubdtype(np.dtype(dtype).type, np.integer), 'geometric needs int dtype'","typeGuard":"def is_int_dtype(d) -> bool:\n    import numpy as np\n    return d is None or d is int or np.issubdtype(np.dtype(d).type, np.integer)","tryCatchPattern":null,"preventionTips":["Remember geometric is discrete: default dtype is int","Cast to float downstream, not at sampling"],"tags":["jax","random","dtype-validation","discrete"],"backgroundTag":"invalid-dtype-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}