{"record":{"id":"08b05642416eb7ff","repo":"jax-ml/jax","slug":"only-byte-aligned-shapes-are-supported-got-shape","errorCode":null,"errorMessage":"Only byte-aligned shapes are supported. Got shape: {ref.dtype}{ref.shape}","messagePattern":"Only byte-aligned shapes are supported\\. Got shape: (.+?)(.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/core.py","lineNumber":498,"sourceCode":"\n# A tree of `GPUMemoryRef`s.\n_GPUMemoryRefTree = Any\n\n\ndef _ref_group_size(refs: _GPUMemoryRefTree) -> int:\n  size = 0\n  for ref in jax.tree.leaves(refs):\n    # Make sure that the start of each ref is aligned with `SMEM_ALIGNMENT`.\n    size = align_to(size, SMEM_ALIGNMENT)\n    if jnp.issubdtype(ref.dtype, jnp.integer):\n      nbits = jnp.iinfo(ref.dtype).bits\n    elif jnp.issubdtype(ref.dtype, jnp.floating):\n      nbits = jnp.finfo(ref.dtype).bits\n    else:\n      raise NotImplementedError(f\"Unsupported dtype: {ref.dtype}\")\n    ref_bits = math.prod(ref.shape) * nbits\n    if ref_bits % 8:\n      raise ValueError(\n          \"Only byte-aligned shapes are supported. Got shape:\"\n          f\" {ref.dtype}{ref.shape}\"\n      )\n    size += ref_bits // 8\n  return size\n\n\ndef _ref_group_tmem_col_size(refs: _GPUMemoryRefTree) -> int:\n  \"\"\"Returns the total number of TMEM columns used by a group of aliased Refs.\n  \"\"\"\n  ncols = 0\n  for ref in jax.tree.leaves(refs):\n    ref_ncols = ref.layout.cols_in_shape(ref.shape,\n                                         dtypes.itemsize_bits(ref.dtype))\n    ncols += align_to(ref_ncols, TMEM_COL_ALIGNMENT)\n  return ncols\n\n","sourceCodeStart":480,"sourceCodeEnd":516,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/core.py#L480-L516","documentation":"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.","triggerScenarios":"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).","commonSituations":"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.","solutions":["Pad the shape so total elements × bits-per-element is a multiple of 8 bytes (e.g. allocate at least one full byte/word)","Use a wider dtype (int8/float32) instead of sub-byte storage","Re-check the shape arithmetic — often an unintended () or (1,) shape from tree_map over a scalar"],"exampleFix":"# before\nscratch = pl_core.SMEM((), jnp.float32)  # or a shape yielding non-byte-aligned total\n# after\nscratch = pl_core.SMEM((8,), jnp.float32)  # ensure prod(shape)*bits % 8 == 0","handlingStrategy":"validation","validationCode":"import math, jax.numpy as jnp\ndef is_byte_aligned(shape, dtype):\n    nbits = (jnp.iinfo(dtype) if jnp.issubdtype(dtype, jnp.integer) else jnp.finfo(dtype)).bits\n    return math.prod(shape) * nbits % 8 == 0","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Check prod(shape)*bits % 8 == 0 before declaring scratch buffers","Pad tiny buffers up to a full word"],"tags":["jax","pallas","mosaic-gpu","shared-memory","alignment","shape-validation"],"backgroundTag":"memory-alignment-error","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}