{"record":{"id":"0fb79d696d8e480b","repo":"jax-ml/jax","slug":"bitcast-1d-ref-with-bitwidth-change-is-not-support-0fb79d","errorCode":null,"errorMessage":"Bitcast 1D ref with bitwidth change is not supported.","messagePattern":"Bitcast 1D ref with bitwidth change is not supported\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/state/utils.py","lineNumber":89,"sourceCode":"        c if is_ref else ref_get(c, ())\n        for is_ref, c in zip(is_const_ref, all_consts)\n    ]\n    return core.eval_jaxpr(jaxpr, all_consts, *args0, *args1)\n\n  hoisted_jaxpr, _ = pe.trace_to_jaxpr(\n      _hoist, ft.flatten_args(*in_avals),\n      jaxpr.debug_info.with_unknown_names())\n  assert not hoisted_jaxpr.consts, \"All consts should have been converted to refs\"\n  return hoisted_jaxpr\n\n\ndef bitcast(x, dtype: DTypeLike):\n  x_bitwidth = dtypes.itemsize_bits(x.dtype)\n  y_bitwidth = dtypes.itemsize_bits(dtype)\n  shape = list(x.shape)\n  if x_bitwidth != y_bitwidth:\n    if len(shape) < 2:\n      raise NotImplementedError(\n          \"Bitcast 1D ref with bitwidth change is not supported.\"\n      )\n    # Note: this is only valid on TPU.\n    if shape[-2] * x_bitwidth % y_bitwidth != 0:\n      raise ValueError(\n          \"Expected input and output shapes are the same after multiplying\"\n          \" the second-minor dimension by the bitwidths.\"\n      )\n  shape[-2] = shape[-2] * x_bitwidth // y_bitwidth\n  if x_bitwidth < y_bitwidth:\n    ratio = y_bitwidth // x_bitwidth\n    x = x.reshape(*x.shape[:-2], x.shape[-2] // ratio, ratio, -1).swapaxes(\n        -1, -2\n    )\n  y = lax.bitcast_convert_type(x, dtype)\n  if x_bitwidth > y_bitwidth:\n    y = y.swapaxes(-1, -2).reshape(shape)\n  return y","sourceCodeStart":71,"sourceCodeEnd":107,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/state/utils.py#L71-L107","documentation":"jax._src.state.utils.bitcast raises NotImplementedError when bitcasting between dtypes of different bitwidths on a ref with fewer than 2 dimensions. Width-changing bitcasts on 1D refs can't be expressed because there's no second-minor dimension to absorb the size change.","triggerScenarios":"bitcast(ref_1d, dtype) where itemsize_bits differ, e.g. bitcasting a 1D float32 ref to float16 or int8.","commonSituations":"Viewing 1D byte buffers as a different-width dtype (common in TPU packing/serialization code).","solutions":["Reshape the ref to 2D first (e.g. (1, n) or (n//ratio, ratio)), bitcast, then reshape back","Use same-bitwidth dtypes (e.g. float32 <-> int32)","Use .get() and jax.numpy reinterpretation on values instead of the ref"],"exampleFix":"# before\ny = bitcast(ref_1d, jnp.int8)  # ref is f32\n# after\nref2 = ref.reshape(1, ref.shape[0])\ny = bitcast(ref2, jnp.int8).reshape(-1)","handlingStrategy":"fallback","validationCode":"import jax.numpy as jnp\nfrom jax._src import dtypes\nif dtypes.itemsize_bits(x.dtype) != dtypes.itemsize_bits(dtype) and len(x.shape) < 2:\n    x = x.reshape(1, x.shape[0]) if x.ndim == 1 else x","typeGuard":null,"tryCatchPattern":"try:\n    y = bitcast(x, dtype)\nexcept NotImplementedError:\n    y = bitcast(x.reshape(1, -1), dtype).reshape(-1)","preventionTips":["Reshape 1D refs to 2D before width-changing bitcasts","Prefer same-width dtype pairs for 1D buffers"],"tags":["jax","bitcast","shape-restriction"],"backgroundTag":"unsupported-bitcast-layout","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}