jax-ml/jax · error · ValueError

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

Error message

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

What it means

jax.random.gamma requires a floating-point dtype because it performs Gamma-distribution math (log/exp transforms, rejection sampling) that only makes sense in float arithmetic. Integer or complex dtypes are rejected with ValueError before the sampler is traced.

Source

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

      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 ``a.shape``.

  See Also:
    loggamma : sample gamma values in log-space, which can provide improved
      accuracy for small values of ``a``.
  """
  key, _ = _check_prng_key("gamma", key)
  if method not in {'exact', 'approximate'}:
    raise ValueError("method argument to `gamma` must be one of "
                     f"{{'exact', 'approximate'}}, got {method!r}")
  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 `gamma` must be a float "
                     f"dtype, got {dtype}")
  if shape is not None:
    shape = core.canonicalize_shape(shape)
  out_sharding = canonicalize_sharding_for_samplers(out_sharding, "gamma", shape)
  if method == 'approximate':
    return maybe_auto_axes(_gamma_approx, out_sharding,
                           shape=shape, dtype=dtype)(key, a)
  return maybe_auto_axes(_gamma, out_sharding, shape=shape, dtype=dtype)(key, a)


def loggamma(key: ArrayLike,
             a: RealArray,
             shape: Shape | None = None,
             dtype: DTypeLikeFloat | None = None,
             *,
             method: str = 'exact',
             out_sharding: NamedSharding | P | None =None) -> Array:
  """Sample log-gamma random values with given shape and float dtype.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass jnp.float32 or jnp.float64 (with jax_enable_x64=True), or omit dtype to get the default float.
2. Validate the dtype with dtypes.issubdtype(dtype, np.floating) before the call in configurable pipelines.
3. Use jax.random.loggamma instead if you need small-alpha accuracy, still with a float dtype.

Example fix

// before
g = jax.random.gamma(key, 3.0, dtype=jnp.int32)

// after
g = jax.random.gamma(key, 3.0, dtype=jnp.float32)

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

def is_float_dtype(dtype) -> bool:
    return dtypes.issubdtype(dtypes.check_and_canonicalize_user_dtype(dtype or float), np.floating)

When it happens

Trigger: Calling jax.random.gamma(key, a, dtype=jnp.int32) or dtype=jnp.complex64; passing a canonicalized custom dtype that resolves to a non-float type.

Common situations: Reusing one dtype variable across many samplers some of which are int-only (poisson); porting NumPy code where numpy.random.gamma had no dtype argument; f64 support toggles where a bfloat16/float16 choice fails upstream validation elsewhere and int is substituted during debugging.

Related errors

• dtype argument to `exponential` must be a float dtype, got {
• method argument to `gamma` must be one of {'exact', 'approxi
• 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

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