jax-ml/jax · error · ValueError

Arguments to rng_uniform must have identical dtypes, got {}

Error message

Arguments to rng_uniform must have identical dtypes, got {} and {}.

What it means

rng_uniform requires its lower bound a and upper bound b to have exactly the same dtype, since the primitive does no implicit promotion. A dtype mismatch is a ValueError at abstract-eval time.

Source

Thrown at jax/_src/lax/lax.py:9213

def rng_uniform(a, b, shape):
  """Stateful PRNG generator. Experimental and its use is discouraged.

  Returns uniformly distributed random numbers in the range [a, b). If
  b <= a, then the result is undefined, and different implementations may
  return different results.

  You should use jax.random for most purposes; this function exists only for
  niche use cases with special performance requirements.

  This API may be removed at any time.
  """
  a, b = core.auto_insert_reshard(a, b)
  return rng_uniform_p.bind(a, b, shape=tuple(shape))

def _rng_uniform_abstract_eval(a, b, *, shape):
  if a.dtype != b.dtype:
    raise ValueError(
      "Arguments to rng_uniform must have identical dtypes, got {} "
      "and {}.".format(a.dtype, b.dtype))
  if a.shape != () or b.shape != ():
    raise ValueError(
      "Arguments to rng_uniform must be scalars; got shapes {} and {}."
      .format(a.shape, b.shape))
  return a.update(shape=shape, dtype=a.dtype,
                  weak_type=(a.weak_type and b.weak_type))

rng_uniform_p = Primitive("rng_uniform")
rng_uniform_p.def_impl(partial(dispatch.apply_primitive, rng_uniform_p))
rng_uniform_p.def_abstract_eval(_rng_uniform_abstract_eval)

def _rng_uniform_lowering(ctx, a, b, *, shape):
  aval_out, = ctx.avals_out
  shape = mlir.ir_constant(np.array(aval_out.shape, np.int64))
  return [hlo.rng(a, b, shape, hlo.RngDistributionAttr.get('UNIFORM'))]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast both bounds explicitly: lax.rng_uniform(jnp.asarray(a, dtype), jnp.asarray(b, dtype)).
2. Pick one dtype variable and use it for both bounds and the desired output.
3. Avoid mixing np scalars and jnp arrays; wrap both in jnp.asarray.

Example fix

# before
z = lax.rng_uniform(np.float64(0.0), jnp.float32(1.0), shape=(1000,))
# after
dt = jnp.float32
z = lax.rng_uniform(jnp.asarray(0.0, dt), jnp.asarray(1.0, dt), shape=(1000,))

a = jnp.asarray(a, dtype)
b = jnp.asarray(b, dtype)
z = lax.rng_uniform(a, b, shape=shape)

def same_dtype(a, b):
    return a.dtype == b.dtype

When it happens

Trigger: lax.rng_uniform(jnp.float32(0), 1.0) where the Python scalar is treated as a weak-typed but differently-resolved dtype, or rng_uniform(np.float64(0), jnp.float32(1)).

Common situations: Mixing numpy scalars and jnp arrays as bounds; porting code where one bound came from a config (Python float) and the other from a computed f32 array; after a global dtype change (jax_enable_x64 toggles).

Related errors

• Incorrect output dtype for return value #{i}: Expected: {out
• {} function carry input and carry output must have equal typ
• lower and upper arguments to fori_loop must have equal types
• the first argument to pow must have an inexact dtype (float
• {name} argument type error: {lhs.dtype}, {rhs.dtype}

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