{"record":{"id":"1c84052b5c2bb598","repo":"jax-ml/jax","slug":"data-type-bitcast-is-only-supported-from-i8-to-oth","errorCode":null,"errorMessage":"Data type bitcast is only supported from i8 to other types.","messagePattern":"Data type bitcast is only supported from i8 to other types\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/lowering.py","lineNumber":1480,"sourceCode":"    ref: ir.Value, src_dtype: ir.Type, dst_dtype: ir.Type\n) -> ir.Value:\n  \"\"\"Allows bitcasting a SMEM ref from one element type to another.\n\n  Args:\n    ref: the reference to bitcast.\n    src_dtype: the source element type.\n    dst_dtype: the destination element type.\n\n  Returns:\n    A bitcasted version of `ref` with element type `dst_dtype`.\n\n  Raises:\n    ValueError: if the source ref is not in SMEM.\n  \"\"\"\n  if src_dtype == dst_dtype:\n    return ref\n  if src_dtype != ir.IntegerType.get_signless(8):\n    raise NotImplementedError(\n        \"Data type bitcast is only supported from i8 to other types.\"\n    )\n  ref_ty = ir.MemRefType(ref.type)\n  if not mgpu_utils.is_smem_ref(ref_ty):\n    raise ValueError(f\"Only workgroup memory is supported but got {ref}.\")\n  if len(ref_ty.shape) != 1:\n    raise NotImplementedError(\n        \"Data type bitcast is only supported for 1D arrays.\"\n    )\n  [stride], _ = ref_ty.get_strides_and_offset()\n  if stride != 1:\n    raise ValueError(\n        \"Data type bitcast is only supported for contiguous 1D arrays, but got \"\n        f\"stride={stride}.\"\n    )\n  [shape_bytes] = ref_ty.shape\n  shape_bitwidth = shape_bytes * 8\n  target_bitwidth = mgpu_utils.bitwidth(dst_dtype)","sourceCodeStart":1462,"sourceCodeEnd":1498,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/lowering.py#L1462-L1498","documentation":"Raised by JAX's Mosaic GPU Pallas lowering when a block reference's element type must be reinterpreted (bitcast) to another dtype, but the source dtype is not i8. The lowering path (_handle_dtype_bitcast, reached via _extract_aliased_ref when handling aliased Refs with different dtypes) only supports viewing i8 storage as a wider type. Any other source width (e.g. i32 -> f32) is unimplemented in the compiler backend.","triggerScenarios":"Writing a Pallas/Mosaic GPU kernel where an aliased Ref's dtype differs from the underlying allocation's dtype and the source dtype is not signless i8 — e.g. passing a Ref[f32] that aliases memory typed as i32, or reinterpreting an i16 buffer.","commonSituations":"Porting Triton-style bitwise reinterpret casts to Pallas on GPU; trying to view a typed SMEM buffer as another dtype for packed I/O; version changes where dtype aliasing support is still i8-only.","solutions":["Change the underlying buffer/ref dtype to i8 (jnp.int8 / jnp.uint8) and bitcast to the target type inside the kernel","Do the dtype conversion with explicit loads + jax.lax.bitcast_convert on values instead of aliasing Refs","Perform reinterpretation on the host with .view() before passing buffers to the kernel","Request/await broader bitcast support upstream or implement it in lowering.py"],"exampleFix":"// before\nref_i32 = pallas_utils.ref(..., jnp.int32)\nout = ref_i32.view(jnp.float32)  # source not i8 -> error\n// after\nref_i8 = pallas_utils.ref(..., jnp.uint8)\nout = ref_i8.view(jnp.float32)  # i8 -> f32 is supported","handlingStrategy":"validation","validationCode":"src, dst = jnp.dtype(ref.dtype), jnp.dtype(target_dtype)\nif src != dst and src not in (jnp.int8, jnp.uint8):\n    raise ValueError(f\"bitcast alias requires i8 source, got {src}\")","typeGuard":"def is_bitcastable_alias(src: jnp.dtype, dst: jnp.dtype) -> bool:\n    return src == dst or src in (jnp.int8, jnp.uint8)","tryCatchPattern":"try:\n    view = ref.view(target_dtype)\nexcept NotImplementedError:\n    # fall back to value-level conversion\n    view = ref[...].astype(jnp.uint8)\n","preventionTips":["Store scratch buffers as i8/uint8 when you plan to reinterpret them","Encapsulate view() calls behind a helper that validates the source dtype","Add kernel unit tests covering dtype aliasing paths"],"tags":["jax","pallas","mosaic-gpu","bitcast","dtype","not-implemented"],"backgroundTag":"unsupported-dtype-conversion","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}