jax-ml/jax · error · NotImplementedError

Data type bitcast is only supported from i8 to other types.

Error message

Data type bitcast is only supported from i8 to other types.

What it means

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.

Source

Thrown at jax/_src/pallas/mosaic_gpu/lowering.py:1480

    ref: ir.Value, src_dtype: ir.Type, dst_dtype: ir.Type
) -> ir.Value:
  """Allows bitcasting a SMEM ref from one element type to another.

  Args:
    ref: the reference to bitcast.
    src_dtype: the source element type.
    dst_dtype: the destination element type.

  Returns:
    A bitcasted version of `ref` with element type `dst_dtype`.

  Raises:
    ValueError: if the source ref is not in SMEM.
  """
  if src_dtype == dst_dtype:
    return ref
  if src_dtype != ir.IntegerType.get_signless(8):
    raise NotImplementedError(
        "Data type bitcast is only supported from i8 to other types."
    )
  ref_ty = ir.MemRefType(ref.type)
  if not mgpu_utils.is_smem_ref(ref_ty):
    raise ValueError(f"Only workgroup memory is supported but got {ref}.")
  if len(ref_ty.shape) != 1:
    raise NotImplementedError(
        "Data type bitcast is only supported for 1D arrays."
    )
  [stride], _ = ref_ty.get_strides_and_offset()
  if stride != 1:
    raise ValueError(
        "Data type bitcast is only supported for contiguous 1D arrays, but got "
        f"stride={stride}."
    )
  [shape_bytes] = ref_ty.shape
  shape_bitwidth = shape_bytes * 8
  target_bitwidth = mgpu_utils.bitwidth(dst_dtype)

View on GitHub (pinned to 1e1c6a8fc0)

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

Example fix

// before
ref_i32 = pallas_utils.ref(..., jnp.int32)
out = ref_i32.view(jnp.float32)  # source not i8 -> error
// after
ref_i8 = pallas_utils.ref(..., jnp.uint8)
out = ref_i8.view(jnp.float32)  # i8 -> f32 is supported

Defensive patterns

Strategy: validation

Validate before calling

src, dst = jnp.dtype(ref.dtype), jnp.dtype(target_dtype)
if src != dst and src not in (jnp.int8, jnp.uint8):
    raise ValueError(f"bitcast alias requires i8 source, got {src}")

Type guard

def is_bitcastable_alias(src: jnp.dtype, dst: jnp.dtype) -> bool:
    return src == dst or src in (jnp.int8, jnp.uint8)

Try / catch

try:
    view = ref.view(target_dtype)
except NotImplementedError:
    # fall back to value-level conversion
    view = ref[...].astype(jnp.uint8)

Prevention

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

When it happens

Trigger: 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.

Common situations: 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.

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/1c84052b5c2bb598. Report an issue: GitHub.