{"record":{"id":"74b2c6aaf8243af1","repo":"jax-ml/jax","slug":"threefry-2x32-requires-uint32-arguments-got","errorCode":null,"errorMessage":"threefry_2x32 requires uint32 arguments, got {}","messagePattern":"threefry_2x32 requires uint32 arguments, got (.+?)","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/random/threefry2x32.py","lineNumber":254,"sourceCode":"    platform='oneapi',\n    inline=False)\n\n\n@api.jit(inline=True)\ndef threefry_2x32(keypair, count):\n  \"\"\"Apply the Threefry 2x32 hash.\n\n  Args:\n    keypair: a pair of 32bit unsigned integers used for the key.\n    count: an array of dtype uint32 used for the counts.\n\n  Returns:\n    An array of dtype uint32 with the same shape as `count`.\n  \"\"\"\n  key1, key2 = keypair\n  if not lax.dtype(key1) == lax.dtype(key2) == lax.dtype(count) == np.uint32:\n    msg = \"threefry_2x32 requires uint32 arguments, got {}\"\n    raise TypeError(msg.format([lax.dtype(x) for x in [key1, key2, count]]))\n\n  flat_count = count.ravel()\n  odd_size = flat_count.shape[0] % 2\n  if core.is_constant_dim(odd_size):\n    if odd_size:\n      x = list(jnp.split(jnp.concatenate([flat_count, jnp.uint32([0])]), 2))\n    else:\n      x = list(jnp.split(flat_count, 2))\n  else:\n    # With symbolic shapes we cannot always tell statically if odd_size is true\n    # or false, so we rewrite this without a conditional.\n    flat_count_padded = jnp.concatenate([flat_count, jnp.uint32([0])])\n    flat_count_padded_half_size = flat_count_padded.shape[0] // 2\n    x = [\n      lax_slicing.dynamic_slice(flat_count_padded, (0,),\n                                (flat_count_padded_half_size,)),\n      lax_slicing.dynamic_slice(flat_count_padded,\n                                (flat_count_padded_half_size,),","sourceCodeStart":236,"sourceCodeEnd":272,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/random/threefry2x32.py#L236-L272","documentation":"threefry_2x32 is the low-level 2x32 block cipher at the heart of JAX's default PRNG; its kernel only accepts uint32 operands for both key words and the count/data array. This TypeError fires when any of key1, key2, or count has a different dtype, typically int32/int64 after implicit conversions.","triggerScenarios":"Calling jax._src.random.threefry2x32.threefry_2x32 directly with int64 arrays (x64 mode); passing counts created by jnp.arange (int32 default) without casting; mixing numpy int arrays with uint32 keys.","commonSituations":"X64-enabled environments where literals default to int64; researchers using the low-level threefry API directly for custom hash chains; data loaded from int-typed sources fed as 'count'.","solutions":["Cast all three operands: key1.astype(jnp.uint32), key2.astype(jnp.uint32), count.astype(jnp.uint32)","Prefer the public APIs (jax.random.fold_in, key_data round-trips) which handle dtype canonicalization","Be explicit with dtype in literals: jnp.uint32([...]) instead of jnp.array([...])"],"exampleFix":"// before\nout = threefry_2x32((k1, k2), jnp.arange(8))\n\n// after\nout = threefry_2x32((k1, k2), jnp.arange(8, dtype=jnp.uint32))","handlingStrategy":"validation","validationCode":"import jax.numpy as jnp\nkey1, key2, count = (jnp.asarray(x, dtype=jnp.uint32) for x in (key1, key2, count))","typeGuard":"import jax.numpy as jnp\ndef all_uint32(*xs) -> bool:\n    return all(jnp.asarray(x).dtype == jnp.uint32 for x in xs)","tryCatchPattern":null,"preventionTips":["Cast every operand to uint32 before low-level threefry calls","Write literals as jnp.uint32([...]); prefer public random APIs"],"tags":["jax","prng","threefry","uint32","dtype"],"backgroundTag":"jax-uint32-dtype-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}