jax-ml/jax · error · ValueError

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

Error message

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

What it means

jax.random.exponential only supports floating-point output dtypes because it computes samples via log-transformation of uniform bits. If the user passes an integer or complex dtype (e.g. jnp.int32), the float-domain math is not defined, so JAX rejects it with ValueError before tracing.

Source

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

    dtype: optional, a float dtype for the returned values (default float64 if
      jax_enable_x64 is true, otherwise float32).
    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("exponential", 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 `exponential` must be a float "
                     f"dtype, got {dtype}")
  shape = core.canonicalize_shape(shape)
  out_sharding = canonicalize_sharding_for_samplers(out_sharding, "exponential", shape)
  return maybe_auto_axes(_exponential, out_sharding,
                         shape=shape, dtype=dtype)(key)

@jit(static_argnums=(1, 2))
def _exponential(key, shape, dtype) -> Array:
  _check_shape("exponential", shape)
  u = uniform(key, shape, dtype)
  # taking 1 - u to move the domain of log to (0, 1] instead of [0, 1)
  return lax.neg(lax.log1p(lax.neg(u)))


def _gamma_one(key: Array, alpha, log_space) -> Array:
  # Ref: A simple method for generating gamma variables, George Marsaglia and Wai Wan Tsang
  # The algorithm can also be founded in:
  # https://en.wikipedia.org/wiki/Gamma_distribution#Generating_gamma-distributed_random_variables

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass a float dtype such as jnp.float32 (the default when dtype=None) or remove the dtype argument entirely.
2. If a configurable dtype must be used, validate dtypes.issubdtype(dtype, np.floating) before the call and coerce with float if not.
3. Audit shared dtype constants in your config so integer dtypes are not reused by float-only samplers.

Example fix

// before
x = jax.random.exponential(key, (1000,), dtype=jnp.int32)

// after
x = jax.random.exponential(key, (1000,), dtype=jnp.float32)

import jax.numpy as jnp, numpy as np
from jax._src import dtypes
assert dtypes.issubdtype(jnp.dtype(dtype).type, np.floating), 'exponential needs a float dtype'

def is_float_dtype(dtype) -> bool:
    return dtypes.issubdtype(jnp.dtype(dtype).type, np.floating)

When it happens

Trigger: Calling jax.random.exponential(key, shape, dtype=jnp.int32) or any non-floating dtype such as jnp.complex64, or a custom dtype alias that canonicalizes to an integer type.

Common situations: Copy-pasting dtype from a different sampler (e.g. rademacher or poisson which accept int dtypes) into exponential; a config file storing a single dtype used for many samplers; assuming dtype=None gives int like some older APIs.

Related errors

• dtype argument to `gamma` must be a float dtype, got {dtype}
• dtype argument to `gumbel` must be a float dtype, got {dtype
• dtype argument to `laplace` must be a float dtype, got {dtyp
• dtype argument to `logistic` must be a float dtype, got {dty
• dtype argument to `pareto` must be a float dtype, got {dtype

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