jax-ml/jax · error · ValueError

Arguments to rng_uniform must be scalars; got shapes {} and

Error message

Arguments to rng_uniform must be scalars; got shapes {} and {}.

What it means

rng_uniform's bounds a and b must be scalars (shape ()). Passing arrays of any other shape is rejected because the primitive defines a scalar interval; the output shape comes only from the shape parameter.

Source

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

  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'))]

mlir.register_lowering(rng_uniform_p, _rng_uniform_lowering)


def _rng_bit_generator_shape_rule(key, *, shape, dtype, algorithm, out_sharding):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. For per-element ranges, use jax.random.uniform with a broadcastable min/max, or compute u01 = rng_uniform(0,1,...) and rescale: lo + u01 * (hi - lo).
2. Squeeze bounds to scalars: jnp.asarray(lo).squeeze().
3. Pass true scalars: lax.rng_uniform(0.0, 1.0, shape).

Example fix

# before
z = lax.rng_uniform(lo_arr, hi_arr, shape=(n,))  # lo_arr/hi_arr: (n,)
# after
u = lax.rng_uniform(jnp.float32(0), jnp.float32(1), shape=(n,))
z = lo_arr + u * (hi_arr - lo_arr)

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

def are_scalars(a, b):
    return a.shape == () and b.shape == ()

When it happens

Trigger: lax.rng_uniform(jnp.zeros((3,)), jnp.ones((3,)), shape=(3,)) — trying to get per-element ranges. Also bounds that became 1-element arrays via jnp.asarray([lo]).

Common situations: Wanting batched/vectorized ranges (e.g., uniform samples with different min/max per row) and assuming broadcasting works; bounds produced by slicing that keep a length-1 dimension.

Related errors

• scan got `length` argument of {} which disagrees with leadin
• conv_general_dilated batch_group_count must divide lhs batch
• conv_general_dilated rhs output feature dimension size must
• conv_general_dilated window and window_strides must have the
• Wrong number of explicit pads for convolution: expected {},

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