{"record":{"id":"78048483938f25d2","repo":"jax-ml/jax","slug":"dtype-argument-to-wald-must-be-a-float-dtype-go","errorCode":null,"errorMessage":"dtype argument to `wald` must be a float dtype, got {dtype}","messagePattern":"dtype argument to `wald` must be a float dtype, got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/random/core.py","lineNumber":3265,"sourceCode":"      jax_enable_x64 is true, otherwise float32).\n    out_sharding: optional, specifies how the output array should be sharded\n      across devices in multi-device computation. Can be a\n      :class:`~jax.sharding.NamedSharding`, a :class:`~jax.sharding.PartitionSpec`\n      (``P``), or ``None`` (default). When specified, the output will be sharded\n      according to the given sharding specification. Primarily used in explicit\n      sharding mode.\n      See the `explicit sharding tutorial <https://docs.jax.dev/en/latest/parallel.html>`_\n      for more details.\n\n  Returns:\n    A random array with the specified dtype and with shape given by ``shape`` if\n    ``shape`` is not None, or else by ``mean.shape``.\n  \"\"\"\n  key, _ = _check_prng_key(\"wald\", key)\n  dtype = dtypes.check_and_canonicalize_user_dtype(\n      float if dtype is None else dtype)\n  if not dtypes.issubdtype(dtype, np.floating):\n    raise ValueError(\"dtype argument to `wald` must be a float \"\n                     f\"dtype, got {dtype}\")\n  shape = _check_broadcast_shapes(\"wald\", shape, mean)\n  out_sharding = canonicalize_sharding(out_sharding, \"wald\")\n  _check_all_safe_to_cast(\"wald\", dtype, mean)\n  return maybe_auto_axes(_wald, out_sharding, shape=shape, dtype=dtype)(key, mean)\n\n@jit(static_argnums=(2, 3))\ndef _wald(key, mean, shape, dtype) -> Array:\n  k1, k2 = _split(key, 2)\n  mean = mean.astype(dtype)\n  mean = jnp.broadcast_to(mean, shape)\n  v = normal(k1, shape, dtype)\n  z = uniform(k2, shape, dtype)\n  y = lax.integer_pow(v, 2)\n  y_sq = lax.integer_pow(y, 2)\n  mean_sq = lax.integer_pow(mean, 2)\n  sqrt_term = lax.sqrt(4 * mean * y + mean_sq * y_sq)\n  x = mean + mean_sq * y / 2 - mean / 2 * sqrt_term","sourceCodeStart":3247,"sourceCodeEnd":3283,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/random/core.py#L3247-L3283","documentation":"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.","triggerScenarios":"Calling jax.random.wald(key, mean, shape, dtype) with dtype=np.int64, np.bool_, or any non-float dtype.","commonSituations":"Forwarding a dtype chosen for a different sampler, or deriving dtype from an integer mean array (e.g. mean stored as int).","solutions":["Omit dtype or pass a float dtype such as jnp.float32","Convert integer-valued mean inputs to float: jnp.asarray(mean, jnp.float32)"],"exampleFix":"// before\nw = jax.random.wald(key, mean, dtype=jnp.int32)\n// after\nw = jax.random.wald(key, jnp.asarray(mean, jnp.float32), dtype=jnp.float32)","handlingStrategy":"validation","validationCode":"import numpy as np\nassert dtype is None or np.issubdtype(np.dtype(dtype).type, np.floating), 'wald needs float dtype'","typeGuard":"def is_float_dtype(d) -> bool:\n    import numpy as np\n    return d is None or np.issubdtype(np.dtype(d).type, np.floating)","tryCatchPattern":null,"preventionTips":["Convert integer mean inputs with jnp.asarray(mean, jnp.float32)","Document float-only dtypes in wrapper docstrings"],"tags":["jax","random","dtype-validation"],"backgroundTag":"invalid-dtype-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}