{"record":{"id":"1dd57664f51bc93f","repo":"jax-ml/jax","slug":"invalid-dtype-for-swap-ref-dtype-expected-out-1dd576","errorCode":null,"errorMessage":"Invalid dtype for `swap`. Ref dtype: {expected_out_ty.dtype}. Value dtype: {val_aval.dtype}. ","messagePattern":"Invalid dtype for `swap`\\. Ref dtype: (.+?)\\. Value dtype: (.+?)\\. ","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/state/primitives.py","lineNumber":438,"sourceCode":"        ref_aval, val_aval, *args, tree=tree)\n  out_aval: core.AbstractValue\n  if not isinstance(ref_aval, AbstractRef):\n    raise ValueError(f\"`swap` must be called on `Ref` types: {ref_aval}.\")\n  if isinstance(val_aval, AbstractRef):\n    raise ValueError(\"Cannot store a Ref into another Ref. \"\n                     \"Did you forget to load from it using `[...]`?\")\n  if isinstance(ref_aval.inner_aval, core.ShapedArray):\n    assert isinstance(val_aval, core.ShapedArray)\n    expected_out_ty = transform_type(transforms, ref_aval.inner_aval)\n    assert isinstance(expected_out_ty, core.ShapedArray)\n    if expected_out_ty.shape != val_aval.shape:\n      raise ValueError(\"Invalid shape for `swap`. \"\n                       f\"Ref shape: {ref_aval.shape}. \"\n                       f\"Expected shape: {expected_out_ty.shape}. \"\n                       f\"Value shape: {val_aval.shape}. \"\n                       f\"Transforms: {transforms}. \")\n    if expected_out_ty.dtype != val_aval.dtype:\n      raise ValueError(\n          \"Invalid dtype for `swap`. \"\n          f\"Ref dtype: {expected_out_ty.dtype}. \"\n          f\"Value dtype: {val_aval.dtype}. \"\n      )\n    out_aval = expected_out_ty\n  else:\n    if transforms:\n      raise ValueError(\"Cannot index non-shaped array with nontrivial indices.\")\n    out_aval = ref_aval.inner_aval\n  return (out_aval, {WriteEffect(0)})\nswap_p.def_effectful_abstract_eval(_swap_abstract_eval)\n\n\ndef _addupdate_abstract_eval(ref_aval: AbstractRef,\n                             val_aval: core.AbstractValue,\n                             *args: Any, tree):\n  transforms = tree_util.tree_unflatten(tree, args)\n  if not isinstance(ref_aval, AbstractRef):","sourceCodeStart":420,"sourceCodeEnd":456,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/state/primitives.py#L420-L456","documentation":"Raised during abstract evaluation of the `swap` primitive when the value being swapped into a Ref has a dtype different from the Ref's element dtype (after applying index transforms). JAX's state primitives require exact dtype match between the written value and the reference; no implicit casting is performed. The error names both dtypes so the mismatch is immediately visible.","triggerScenarios":"Calling `ref.swap(...)` (or `swap_p.bind`, or lax.fori_loop bodies / scan carry updates that compile to swap) with a value whose dtype differs from the ref, e.g. writing a float32 into an int32 Ref, or a weakly-typed Python scalar array promoted differently.","commonSituations":"Creating a Ref via `jax.experimental.state.ref` or getters like `Ref(np.zeros(..., dtype=np.int32))` then swapping a float loss/gradient; mixing f32/f64 under enabled x64; converting numpy arrays whose default dtype differs from the initialized buffer.","solutions":["Cast the value to the Ref's dtype before swapping: `x.astype(ref.dtype)` (or `ref_aval.inner_aval.dtype`).","Initialize the Ref with the same dtype as the values you will write (e.g. `jnp.zeros(shape, dtype=x.dtype)`).","Check for accidental weak-typed scalars: wrap Python scalars in `jnp.asarray(x, dtype=...)`.","Under `jax.experimental.enable_x64`, verify both sides weren't created under different x64 settings."],"exampleFix":"// before\nref = state.Ref(jnp.zeros((n,), jnp.int32))\nold = ref.swap(0, jnp.float32(1.5))\n// after\nref = state.Ref(jnp.zeros((n,), jnp.int32))\nold = ref.swap(0, jnp.asarray(1, dtype=jnp.int32))","handlingStrategy":"validation","validationCode":"val = jnp.asarray(val)\nassert val.dtype == ref.aval.inner_aval.dtype, (val.dtype, ref.aval.inner_aval.dtype)\nold = ref.swap(idx, val)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Always create refs and the values written to them from a single shared dtype constant.","Use jnp.asarray(x, dtype=ref.dtype) at write boundaries.","Add dtype assertions in tests covering state updates."],"tags":["jax","dtype-mismatch","state-primitives","swap"],"backgroundTag":"dtype-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}