jax-ml/jax · error · ValueError

Only byte-aligned shapes are supported. Got shape: {ref.dtyp

Error message

Only byte-aligned shapes are supported. Got shape: {ref.dtype}{ref.shape}

What it means

Shared-memory refs in Mosaic GPU must occupy a whole number of bytes. `_ref_group_size` multiplies the element count by the dtype's bit width and raises this ValueError if total bits aren't divisible by 8 — e.g. single-element or oddly-shaped sub-byte-type buffers.

Source

Thrown at jax/_src/pallas/mosaic_gpu/core.py:498

# A tree of `GPUMemoryRef`s.
_GPUMemoryRefTree = Any


def _ref_group_size(refs: _GPUMemoryRefTree) -> int:
  size = 0
  for ref in jax.tree.leaves(refs):
    # Make sure that the start of each ref is aligned with `SMEM_ALIGNMENT`.
    size = align_to(size, SMEM_ALIGNMENT)
    if jnp.issubdtype(ref.dtype, jnp.integer):
      nbits = jnp.iinfo(ref.dtype).bits
    elif jnp.issubdtype(ref.dtype, jnp.floating):
      nbits = jnp.finfo(ref.dtype).bits
    else:
      raise NotImplementedError(f"Unsupported dtype: {ref.dtype}")
    ref_bits = math.prod(ref.shape) * nbits
    if ref_bits % 8:
      raise ValueError(
          "Only byte-aligned shapes are supported. Got shape:"
          f" {ref.dtype}{ref.shape}"
      )
    size += ref_bits // 8
  return size


def _ref_group_tmem_col_size(refs: _GPUMemoryRefTree) -> int:
  """Returns the total number of TMEM columns used by a group of aliased Refs.
  """
  ncols = 0
  for ref in jax.tree.leaves(refs):
    ref_ncols = ref.layout.cols_in_shape(ref.shape,
                                         dtypes.itemsize_bits(ref.dtype))
    ncols += align_to(ref_ncols, TMEM_COL_ALIGNMENT)
  return ncols

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pad the shape so total elements × bits-per-element is a multiple of 8 bytes (e.g. allocate at least one full byte/word)
2. Use a wider dtype (int8/float32) instead of sub-byte storage
3. Re-check the shape arithmetic — often an unintended () or (1,) shape from tree_map over a scalar

Example fix

# before
scratch = pl_core.SMEM((), jnp.float32)  # or a shape yielding non-byte-aligned total
# after
scratch = pl_core.SMEM((8,), jnp.float32)  # ensure prod(shape)*bits % 8 == 0

import math, jax.numpy as jnp
def is_byte_aligned(shape, dtype):
    nbits = (jnp.iinfo(dtype) if jnp.issubdtype(dtype, jnp.integer) else jnp.finfo(dtype)).bits
    return math.prod(shape) * nbits % 8 == 0

When it happens

Trigger: Allocating an SMEM ref whose total bit count isn't byte-aligned, e.g. shape () or (1,) with a 4-bit-ish dtype, or any shape where prod(shape)*nbits % 8 != 0 (most commonly tiny bool/sub-byte arrays after dtype handling, or zero-size refs with unusual dtypes).

Common situations: Creating scalar or single-element scratch buffers; shrinking a validated buffer shape during tuning; using packed/nibble representations that JAX's storage doesn't actually support for SMEM.

Related errors

• No valid out swizzle{what}: minor dimension has {minor_dim_b
• Unsupported dtype: {ref.dtype}
• Can not bitcast memory region of size {shape_bitwidth} bits
• Only byte-aligned bitcasts are supported.
• Can only transfer integer bytes (shape={shape}, dtype={dtype

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