jax-ml/jax · error · ValueError

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

Error message

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

What it means

jax.random.wald requires a floating-point dtype. The dtype is canonicalized (defaulting to float) and validated with dtypes.issubdtype(dtype, np.floating); integer, bool, or complex dtypes raise this ValueError.

Source

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

      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 dtype and with shape given by ``shape`` if
    ``shape`` is not None, or else by ``mean.shape``.
  """
  key, _ = _check_prng_key("wald", key)
  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 `wald` must be a float "
                     f"dtype, got {dtype}")
  shape = _check_broadcast_shapes("wald", shape, mean)
  out_sharding = canonicalize_sharding(out_sharding, "wald")
  _check_all_safe_to_cast("wald", dtype, mean)
  return maybe_auto_axes(_wald, out_sharding, shape=shape, dtype=dtype)(key, mean)

@jit(static_argnums=(2, 3))
def _wald(key, mean, shape, dtype) -> Array:
  k1, k2 = _split(key, 2)
  mean = mean.astype(dtype)
  mean = jnp.broadcast_to(mean, shape)
  v = normal(k1, shape, dtype)
  z = uniform(k2, shape, dtype)
  y = lax.integer_pow(v, 2)
  y_sq = lax.integer_pow(y, 2)
  mean_sq = lax.integer_pow(mean, 2)
  sqrt_term = lax.sqrt(4 * mean * y + mean_sq * y_sq)
  x = mean + mean_sq * y / 2 - mean / 2 * sqrt_term

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Omit dtype or pass a float dtype such as jnp.float32
2. Convert integer-valued mean inputs to float: jnp.asarray(mean, jnp.float32)

Example fix

// before
w = jax.random.wald(key, mean, dtype=jnp.int32)
// after
w = jax.random.wald(key, jnp.asarray(mean, jnp.float32), dtype=jnp.float32)

import numpy as np
assert dtype is None or np.issubdtype(np.dtype(dtype).type, np.floating), 'wald needs float dtype'

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

When it happens

Trigger: Calling jax.random.wald(key, mean, shape, dtype) with dtype=np.int64, np.bool_, or any non-float dtype.

Common situations: Forwarding a dtype chosen for a different sampler, or deriving dtype from an integer mean array (e.g. mean stored as int).

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 `geometric` must be an int dtype, got {dty
• 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/78048483938f25d2. Report an issue: GitHub.