jax-ml/jax · error · ValueError

dtype argument to `chisquare` must be a float dtype, got {dt

Error message

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

What it means

jax.random.chisquare requires a floating-point dtype because it samples via gamma variates in float arithmetic. Integer or complex dtypes raise ValueError. Subsequent checks also require shape to broadcast against df.shape and df to be safely castable to dtype.

Source

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

      :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 ``df.shape``.
  """
  key, _ = _check_prng_key("chisquare", key)
  if method not in {"exact", "approximate"}:
    raise ValueError("method argument to `chisquare` 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("dtype argument to `chisquare` must be a float "
                     f"dtype, got {dtype}")
  shape = _check_broadcast_shapes("chisquare", shape, df)
  _check_all_safe_to_cast("chisquare", dtype, df)
  out_sharding = canonicalize_sharding_for_samplers(out_sharding, "chisquare", shape)
  return maybe_auto_axes(_chisquare, out_sharding, method=method,
                         shape=shape, dtype=dtype)(key, df)


@jit(static_argnums=(2, 3, 4))
def _chisquare(key, df, method, shape, dtype) -> Array:
  df = lax.convert_element_type(df, dtype)
  two = lax._const(df, 2)
  half_df = lax.div(df, two)
  log_g = loggamma(key, a=half_df, shape=shape, dtype=dtype, method=method)
  chi2 = lax.mul(jnp.exp(log_g), two)
  return chi2

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass jnp.float32/jnp.float64 or omit dtype.
2. Check that shape broadcasts against df.shape and that df casts safely to the chosen dtype (e.g. avoid float32 dtype with float64 df under x64).
3. Validate configurable dtypes against np.floating.

Example fix

// before
x = jax.random.chisquare(key, 3.0, dtype=jnp.int32)

// after
x = jax.random.chisquare(key, 3.0, dtype=jnp.float32)

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

def is_float_dtype(dtype) -> bool:
    from jax._src import dtypes
    import numpy as np
    return dtypes.issubdtype(dtypes.check_and_canonicalize_user_dtype(dtype or float), np.floating)

When it happens

Trigger: jax.random.chisquare(key, df, dtype=jnp.int32) or any dtype where dtypes.issubdtype(dtype, np.floating) is False.

Common situations: Chi-square-like count data tempting int dtypes; shared dtype configs; porting numpy.random.chisquare which has no dtype parameter.

Related errors

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

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