jax-ml/jax · error · ValueError

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

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.

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 `chisquare` must be a float dtype, got
• dtype argument to `exponential` must be a float dtype, got
• dtype argument to `f` must be a float dtype, got
• dtype argument to `gumbel` must be a float dtype, got
• dtype argument to `laplace` must be a float dtype, got

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

Appendix: 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)