{"record":{"id":"8f6df35bc155845f","repo":"jax-ml/jax","slug":"arguments-to-rng-uniform-must-have-identical-dtype","errorCode":null,"errorMessage":"Arguments to rng_uniform must have identical dtypes, got {} and {}.","messagePattern":"Arguments to rng_uniform must have identical dtypes, got (.+?) and (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":9213,"sourceCode":"\ndef rng_uniform(a, b, shape):\n  \"\"\"Stateful PRNG generator. Experimental and its use is discouraged.\n\n  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","sourceCodeStart":9195,"sourceCodeEnd":9231,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L9195-L9231","documentation":"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.","triggerScenarios":"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)).","commonSituations":"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).","solutions":["Cast both bounds explicitly: lax.rng_uniform(jnp.asarray(a, dtype), jnp.asarray(b, dtype)).","Pick one dtype variable and use it for both bounds and the desired output.","Avoid mixing np scalars and jnp arrays; wrap both in jnp.asarray."],"exampleFix":"# before\nz = lax.rng_uniform(np.float64(0.0), jnp.float32(1.0), shape=(1000,))\n# after\ndt = jnp.float32\nz = lax.rng_uniform(jnp.asarray(0.0, dt), jnp.asarray(1.0, dt), shape=(1000,))","handlingStrategy":"validation","validationCode":"a = jnp.asarray(a, dtype)\nb = jnp.asarray(b, dtype)\nz = lax.rng_uniform(a, b, shape=shape)","typeGuard":"def same_dtype(a, b):\n    return a.dtype == b.dtype","tryCatchPattern":null,"preventionTips":["Wrap both bounds in jnp.asarray with an explicit shared dtype.","Don't mix np scalars and jnp arrays as bounds."],"tags":["jax","rng-uniform","dtype-mismatch"],"backgroundTag":"dtype-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}