{"record":{"id":"d251a9ce6116521d","repo":"jax-ml/jax","slug":"expected-input-and-output-shapes-are-the-same-afte","errorCode":null,"errorMessage":"Expected input and output shapes are the same after multiplying the second-minor dimension by the bitwidths.","messagePattern":"Expected input and output shapes are the same after multiplying the second-minor dimension by the bitwidths\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/state/utils.py","lineNumber":94,"sourceCode":"  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\n\n\ndef eval_bitcast_shape(x, dtype: DTypeLike):\n  f = partial(bitcast, dtype=dtype)\n  return api.eval_shape(f, api.ShapeDtypeStruct.like(x)).shape","sourceCodeStart":76,"sourceCodeEnd":112,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/state/utils.py#L76-L112","documentation":"jax._src.state.utils.bitcast raises ValueError when the second-minor dimension's total bits (shape[-2] * x_bitwidth) aren't evenly divisible by the target bitwidth, so no valid output shape exists for the width-changing bitcast. The comment notes this packing scheme is TPU-specific.","triggerScenarios":"bitcast(x, dtype) with differing bitwidths where shape[-2] * old_bits % new_bits != 0, e.g. a (5, 4) f32 ref bitcast to f16: 5*32=160 not divisible by 16... (concretely, dims whose bit total doesn't factor the new width).","commonSituations":"Bitcasting to wider/narrower dtypes on TPU without padding dimensions to a multiple of the width ratio.","solutions":["Pad or trim shape[-2] so shape[-2]*old_bits is divisible by new_bits","Reshape so the second-minor dim is a multiple of (new_bits//old_bits) or vice versa","Use same-bitwidth dtypes to avoid the constraint"],"exampleFix":"# before\ny = bitcast(x_f32_shape_5x4, jnp.f16)\n# after\nx = x.reshape(10, 2)  # make second-minor dim compatible\ny = bitcast(x, jnp.f16)","handlingStrategy":"validation","validationCode":"from jax._src import dtypes\nxb, yb = dtypes.itemsize_bits(x.dtype), dtypes.itemsize_bits(dtype)\nassert xb == yb or (len(x.shape) >= 2 and x.shape[-2] * xb % yb == 0), \"incompatible bitcast shape\"","typeGuard":null,"tryCatchPattern":"try:\n    y = bitcast(x, dtype)\nexcept ValueError as e:\n    if \"bitwidths\" in str(e):\n        pad = (-x.shape[-2]) % (dtypes.itemsize_bits(dtype) // dtypes.itemsize_bits(x.dtype))\n        x = x.reshape(x.shape[-2] + pad, -1)\n        y = bitcast(x, dtype)\n    else:\n        raise","preventionTips":["Pad second-minor dims to width-ratio multiples before bitcast","Validate divisibility in a helper before calling bitcast"],"tags":["jax","bitcast","divisibility"],"backgroundTag":"shape-incompatible-bitcast","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}