jax-ml/jax · error · ValueError

Data type bitcast is only supported for contiguous 1D arrays

Error message

Data type bitcast is only supported for contiguous 1D arrays, but got stride={stride}.

What it means

Bitcasting bytes to a wider dtype requires physically contiguous memory, because the lowering produces a new memref with stride 1 in the target dtype. If the source i8 memref has a stride other than 1 (e.g. a strided slice of SMEM), the reinterpretation would be incorrect and is rejected.

Source

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

  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)

  if shape_bitwidth % target_bitwidth:
    raise ValueError(
        f"Can not bitcast memory region of size {shape_bitwidth} bits to dtype "
        f"with {target_bitwidth} bits."
    )

  result_type = ir.MemRefType.get(
      shape=(shape_bitwidth // target_bitwidth,),
      element_type=dst_dtype,
      memory_space=ref_ty.memory_space,
  )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Materialize a contiguous copy of the data before bitcasting
Adjust slicing so the resulting memref is contiguous (slice the innermost dim fully)
Use gather/scatter loads instead of memory reinterpretation

Example fix

// before
s = buf[::2]           # stride 2
v = s.view(jnp.float32)
// after
s = buf[:n]            # contiguous
v = s.view(jnp.float32)

Defensive patterns

Strategy: validation

Validate before calling

import numpy as np
assert all(s == 0 or np.prod(buf.shape[max(i+1,0):]) == buf.shape[-1] for i, s in enumerate(getattr(buf, 'strides', (1,)))), 'ref must be contiguous before view'

Type guard

def is_contiguous_1d(buf) -> bool:
    return len(buf.shape) == 1 and (getattr(buf, 'strides', (1,))[0] in (1, None))

Prevention

Avoid strided slices of buffers you intend to bitcast
Copy into a fresh contiguous buffer when in doubt
Keep reinterpretation on the innermost contiguous axis

When it happens

Trigger: Aliasing a non-contiguous (strided) i8 SMEM ref with a different dtype — e.g. taking a column slice or every-other-element view then calling .view(dtype).

Common situations: Slicing scratch buffers before reinterpretation; strided views produced by block transformations during lowering of aliased Refs.

Related errors

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