jax-ml/jax · error · ValueError

Not implemented: the 2nd minor dim can not be perfectly pack

Error message

Not implemented: the 2nd minor dim can not be perfectly packed or unpacked

What it means

bitcast in Mosaic reinterprets memory by packing src-bitwidth elements of the second-to-minor dimension into dst-bitwidth elements. If x.shape[-2] * src_bitwidth is not divisible by dst_bitwidth, the pack/unpack cannot be done in whole elements, so the primitive rejects it.

Source

Thrown at jax/_src/pallas/mosaic/primitives.py:73


def repeat(x: jax.Array, repeats: int, axis: int) -> jax.Array:
  axis = util.canonicalize_axis(axis, x.ndim)
  reps = [repeats if i == axis else 1 for i in range(x.ndim)]
  return jnp.tile(x, reps)


bitcast_p = jax_core.Primitive("bitcast")


def bitcast(x: jax.Array, ty: DTypeLike) -> jax.Array:
  ty = dtypes.check_and_canonicalize_user_dtype(ty)
  if len(x.shape) < 2:
    raise ValueError("Not implemented: bitcast 1D")
  src_bitwidth = dtypes.itemsize_bits(x.dtype)
  dst_bitwidth = dtypes.itemsize_bits(ty)
  if x.shape[-2] * src_bitwidth % dst_bitwidth:
    raise ValueError(
        "Not implemented: the 2nd minor dim can not be perfectly packed or"
        " unpacked"
    )
  return bitcast_p.bind(x, ty=ty)


@bitcast_p.def_abstract_eval
def _bitcast_abstract_eval(x, *, ty):
  shape = list(x.shape)
  src_bitwidth = dtypes.itemsize_bits(x.dtype)
  dst_bitwidth = dtypes.itemsize_bits(ty)
  shape[-2] = shape[-2] * src_bitwidth // dst_bitwidth
  return jax_core.ShapedArray(shape, ty)


def _bitcast_lowering_rule(ctx: mlir.LoweringRuleContext, x, *, ty):
  def _bitcast(x):
    src_bitwidth = dtypes.itemsize_bits(x.dtype)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Adjust the block/shape so the second-to-minor dimension times source bitwidth divides evenly by destination bitwidth (e.g. make shape[-2] a multiple of dst_bitwidth/src_bitwidth)
Pad the dimension to the next valid size before bitcasting and slice afterwards
Pick a destination dtype whose bitwidth divides the source packing (int32 <-> f32, int8 packs of 4 into int32, etc.)

Example fix

# before
y = mosaic.bitcast(x, jnp.float8_e4m3)  # x.shape == (16, 6) f32 -> 6*32=192 % 8 ok, but e.g. (16,5): 160%8==0; failing case: sub-byte dst

# after
# ensure shape[-2] packs evenly: pad to multiple of dst_bits/src_bits
pad = (-x.shape[-2]) % (dst_bits // src_bits)
x_p = jnp.pad(x, [(0,0),(0,pad)])
y = mosaic.bitcast(x_p, jnp.float8_e4m3)

Defensive patterns

Strategy: validation

Validate before calling

from jax import dtypes
def packs_evenly(x, ty):
  return (x.shape[-2] * dtypes.itemsize_bits(x.dtype)) % dtypes.itemsize_bits(ty) == 0
assert packs_evenly(x, jnp.float8_e4m3), "2nd-minor dim cannot be packed evenly"

Prevention

Choose block shapes whose 2nd-minor dim is a multiple of the bitwidth ratio
Add a shape assertion in kernel wrappers when using sub-byte dtypes

When it happens

Trigger: bitcast(x, ty) where (x.shape[-2] * itemsize_bits(x.dtype)) % itemsize_bits(ty) != 0, e.g. bitcasting a (8, 5) f32 block to f8e4m3 (8*5*32=1280 bits not divisible cleanly per element grouping), or packing sub-byte types where the 2nd-minor dim is not a multiple of the ratio.

Common situations: Quantizing to 8-bit or 4-bit types inside a TPU Pallas kernel with a block whose second-to-minor dimension was chosen for compute tiling rather than bit-packing; switching a kernel from f32 to sub-byte dtypes without adjusting block shapes.

Related errors

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