jax-ml/jax · error · ValueError

dtype argument to `geometric` must be an int dtype, got {dty

Error message

dtype argument to `geometric` must be an int dtype, got {dtype}

What it means

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.

Source

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

      jax_enable_x64 is true, otherwise int32).
    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 dtype and with shape given by ``shape`` if
    ``shape`` is not None, or else by ``p.shape``.
  """
  key, _ = _check_prng_key("geometric", key)
  dtype = dtypes.check_and_canonicalize_user_dtype(
      int if dtype is None else dtype)
  if not dtypes.issubdtype(dtype, np.integer):
    raise ValueError("dtype argument to `geometric` must be an int "
                     f"dtype, got {dtype}")
  shape = _check_broadcast_shapes("geometric", shape, p)
  out_sharding = canonicalize_sharding_for_samplers(out_sharding, "geometric", shape)
  return _geometric(key, p, shape, dtype, out_sharding)

@jit(static_argnums=(2, 3, 4))
def _geometric(key, p, shape, dtype, out_sharding) -> Array:
  check_arraylike("geometric", p)
  p, = promote_dtypes_inexact(p)
  u = uniform(key, shape, p.dtype, out_sharding=out_sharding)
  # TODO(jakevdp): switch to log_u = lax.log1p(u - 1)
  # For now we map u=0 to u=1 to avoid inf in log_u without otherwise
  # changing samples produced for a given key.
  u = jnp.where(u == 0, 1, u)
  log_u = lax.log(u)
  log_one_minus_p = lax.log1p(-p)
  log_one_minus_p = jnp.broadcast_to(log_one_minus_p, shape, out_sharding=out_sharding)
  g = lax.floor(lax.div(log_u, log_one_minus_p)) + 1

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass an int dtype such as np.int32 or jnp.uint32, or omit dtype
2. If floats are needed downstream, sample as int then cast: jnp.asarray(x, np.float32)

Example fix

// before
g = jax.random.geometric(key, 0.3, dtype=jnp.float32)
// after
g = jax.random.geometric(key, 0.3, dtype=jnp.int32)
# if floats needed:
gf = jax.random.geometric(key, 0.3, dtype=jnp.int32).astype(jnp.float32)

import numpy as np
assert dtype is None or np.issubdtype(np.dtype(dtype).type, np.integer), 'geometric needs int dtype'

def is_int_dtype(d) -> bool:
    import numpy as np
    return d is None or d is int or np.issubdtype(np.dtype(d).type, np.integer)

When it happens

Trigger: Calling jax.random.geometric(key, p, shape, dtype=np.float32) or any non-integer dtype.

Common situations: Using a float dtype out of habit from continuous samplers; np.bool_ is not np.integer so boolean requests also fail.

Related errors

• dtype argument to `double_sided_maxwell` must be a float dty
• dtype argument to `weibull_min` must be a float dtype, got {
• dtype argument to `rayleigh` must be a float dtype, got {dty
• dtype argument to `wald` must be a float dtype, got {dtype}
• dtype argument to `triangular` must be a float dtype, got {d

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