{"record":{"id":"afe6a5f7d592406f","repo":"jax-ml/jax","slug":"bitcast-convert-type-with-different-bitwidths-not","errorCode":null,"errorMessage":"bitcast_convert_type with different bitwidths not supported yet: {old_dtype=}, {new_dtype=}","messagePattern":"bitcast_convert_type with different bitwidths not supported yet: (.+?), (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/fuser/block_spec.py","lineNumber":2205,"sourceCode":"  return [block_transform]\n\n\n@register_eval_rule(lax.bitcast_convert_type_p)\ndef _bitcast_convert_type_eval_rule(eval_ctx: KernelEvalContext, x, new_dtype):\n  del eval_ctx\n  return jax.lax.bitcast_convert_type(x, new_dtype)\n\n\n@register_pull_block_spec_rule(lax.bitcast_convert_type_p)\ndef _bitcast_convert_type_pull_rule(\n    ctx: PullRuleContext,\n    block_transform: BlockIndexTransform,\n    *,\n    new_dtype: jnp.dtype,\n):\n  old_dtype = ctx.avals_in[0].dtype\n  if old_dtype.itemsize != new_dtype.itemsize:\n    raise NotImplementedError(\n        'bitcast_convert_type with different bitwidths not supported yet:'\n        f' {old_dtype=}, {new_dtype=}'\n    )\n  return [block_transform]\n\n\n@register_eval_rule(prng.random_bits_p)\ndef _random_bits_eval_rule(eval_ctx: KernelEvalContext, key, bit_width, shape):\n  del shape\n  block_spec = eval_ctx.out_block_specs[0]\n  indices = eval_ctx.get_out_block_indices()[0]\n  block_shape = block_spec.block_shape\n  # This is the important part here: we fold in block indices into the key so\n  # each block gets different random numbers.\n  for idx in indices:\n    key = jax.random.fold_in(key, idx)\n  return prng.random_bits(key, bit_width=bit_width, shape=block_shape)\n","sourceCodeStart":2187,"sourceCodeEnd":2223,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/fuser/block_spec.py#L2187-L2223","documentation":"The fuser only supports bitcast_convert_type when the source and destination dtypes have identical bit widths (itemsize). Bitcasting across widths (e.g. float32 -> float16, or int32 -> int8) changes the array extent per block, which block specs can't yet express, hence NotImplementedError.","triggerScenarios":"lax.bitcast_convert_type inside a fused Pallas region where old_dtype.itemsize != new_dtype.itemsize, e.g. bitcasting a f32 array to uint16, or float16 to int8.","commonSituations":"Packing/unpacking sub-word types (fp8/bf16) for TPU/GPU kernels; converting a wide accumulator to a narrow storage dtype via bitcast instead of a value-preserving cast; assuming bitcast behaves like astype.","solutions":["Use equal-width bitcasts (f32<->u32, f16<->u16, etc.) and do any width change with a real cast (lax.convert_element_type) instead","Perform the width-changing bitcast outside the fused kernel and pass the result as input","If narrowing is required, reshape/atomically split via supported ops before bitcasting equal-width pieces"],"exampleFix":"// before\ny = lax.bitcast_convert_type(x, jnp.dtype('float16'))  # x is float32: widths differ\n// after\ny = lax.bitcast_convert_type(x, jnp.dtype('uint32'))  # equal width, then handle narrowing separately\n","handlingStrategy":"validation","validationCode":"assert x.dtype.itemsize == jnp.dtype(new_dtype).itemsize, 'bitcast requires equal itemsize'","typeGuard":"def bitcast_widths_match(old: jnp.dtype, new: jnp.dtype) -> bool:\n    return old.itemsize == new.itemsize","tryCatchPattern":"try:\n    y = fused_bitcast(x, new_dtype)\nexcept NotImplementedError as e:\n    if 'different bitwidths' in str(e):\n        y = lax.bitcast_convert_type(x, new_dtype)  # outside fusion\n    else:\n        raise","preventionTips":["Bitcast only equal-width types (f32<->u32, f16<->u16)","Use convert_element_type for width changes","Centralize dtype-pair checks in a helper"],"tags":["jax","pallas","bitcast","dtype","not-implemented"],"backgroundTag":"dtype-bitwidth-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}