{"record":{"id":"9bde85f3597bb831","repo":"jax-ml/jax","slug":"arguments-to-rng-uniform-must-be-scalars-got-shap","errorCode":null,"errorMessage":"Arguments to rng_uniform must be scalars; got shapes {} and {}.","messagePattern":"Arguments to rng_uniform must be scalars; got shapes (.+?) and (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":9217,"sourceCode":"  Returns uniformly distributed random numbers in the range [a, b). If\n  b <= a, then the result is undefined, and different implementations may\n  return different results.\n\n  You should use jax.random for most purposes; this function exists only for\n  niche use cases with special performance requirements.\n\n  This API may be removed at any time.\n  \"\"\"\n  a, b = core.auto_insert_reshard(a, b)\n  return rng_uniform_p.bind(a, b, shape=tuple(shape))\n\ndef _rng_uniform_abstract_eval(a, b, *, shape):\n  if a.dtype != b.dtype:\n    raise ValueError(\n      \"Arguments to rng_uniform must have identical dtypes, got {} \"\n      \"and {}.\".format(a.dtype, b.dtype))\n  if a.shape != () or b.shape != ():\n    raise ValueError(\n      \"Arguments to rng_uniform must be scalars; got shapes {} and {}.\"\n      .format(a.shape, b.shape))\n  return a.update(shape=shape, dtype=a.dtype,\n                  weak_type=(a.weak_type and b.weak_type))\n\nrng_uniform_p = Primitive(\"rng_uniform\")\nrng_uniform_p.def_impl(partial(dispatch.apply_primitive, rng_uniform_p))\nrng_uniform_p.def_abstract_eval(_rng_uniform_abstract_eval)\n\ndef _rng_uniform_lowering(ctx, a, b, *, shape):\n  aval_out, = ctx.avals_out\n  shape = mlir.ir_constant(np.array(aval_out.shape, np.int64))\n  return [hlo.rng(a, b, shape, hlo.RngDistributionAttr.get('UNIFORM'))]\n\nmlir.register_lowering(rng_uniform_p, _rng_uniform_lowering)\n\n\ndef _rng_bit_generator_shape_rule(key, *, shape, dtype, algorithm, out_sharding):","sourceCodeStart":9199,"sourceCodeEnd":9235,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L9199-L9235","documentation":"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.","triggerScenarios":"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]).","commonSituations":"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.","solutions":["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).","Squeeze bounds to scalars: jnp.asarray(lo).squeeze().","Pass true scalars: lax.rng_uniform(0.0, 1.0, shape)."],"exampleFix":"# before\nz = lax.rng_uniform(lo_arr, hi_arr, shape=(n,))  # lo_arr/hi_arr: (n,)\n# after\nu = lax.rng_uniform(jnp.float32(0), jnp.float32(1), shape=(n,))\nz = lo_arr + u * (hi_arr - lo_arr)","handlingStrategy":"validation","validationCode":"a = jnp.asarray(a).squeeze()\nb = jnp.asarray(b).squeeze()\nassert a.shape == () == b.shape\nz = lax.rng_uniform(a, b, shape=shape)","typeGuard":"def are_scalars(a, b):\n    return a.shape == () and b.shape == ()","tryCatchPattern":null,"preventionTips":["For per-element ranges, generate U(0,1) then rescale: lo + u*(hi-lo).","Squeeze bound arrays to scalars before calling rng_uniform."],"tags":["jax","rng-uniform","scalar-requirement","shape-validation"],"backgroundTag":"non-scalar-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}