{"record":{"id":"5f601a37bb9abd6e","repo":"jax-ml/jax","slug":"not-implemented-the-2nd-minor-dim-can-not-be-perf","errorCode":null,"errorMessage":"Not implemented: the 2nd minor dim can not be perfectly packed or unpacked","messagePattern":"Not implemented: the 2nd minor dim can not be perfectly packed or unpacked","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/primitives.py","lineNumber":73,"sourceCode":"\n\ndef repeat(x: jax.Array, repeats: int, axis: int) -> jax.Array:\n  axis = util.canonicalize_axis(axis, x.ndim)\n  reps = [repeats if i == axis else 1 for i in range(x.ndim)]\n  return jnp.tile(x, reps)\n\n\nbitcast_p = jax_core.Primitive(\"bitcast\")\n\n\ndef bitcast(x: jax.Array, ty: DTypeLike) -> jax.Array:\n  ty = dtypes.check_and_canonicalize_user_dtype(ty)\n  if len(x.shape) < 2:\n    raise ValueError(\"Not implemented: bitcast 1D\")\n  src_bitwidth = dtypes.itemsize_bits(x.dtype)\n  dst_bitwidth = dtypes.itemsize_bits(ty)\n  if x.shape[-2] * src_bitwidth % dst_bitwidth:\n    raise ValueError(\n        \"Not implemented: the 2nd minor dim can not be perfectly packed or\"\n        \" unpacked\"\n    )\n  return bitcast_p.bind(x, ty=ty)\n\n\n@bitcast_p.def_abstract_eval\ndef _bitcast_abstract_eval(x, *, ty):\n  shape = list(x.shape)\n  src_bitwidth = dtypes.itemsize_bits(x.dtype)\n  dst_bitwidth = dtypes.itemsize_bits(ty)\n  shape[-2] = shape[-2] * src_bitwidth // dst_bitwidth\n  return jax_core.ShapedArray(shape, ty)\n\n\ndef _bitcast_lowering_rule(ctx: mlir.LoweringRuleContext, x, *, ty):\n  def _bitcast(x):\n    src_bitwidth = dtypes.itemsize_bits(x.dtype)","sourceCodeStart":55,"sourceCodeEnd":91,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/primitives.py#L55-L91","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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.)"],"exampleFix":"# before\ny = 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\n\n# after\n# ensure shape[-2] packs evenly: pad to multiple of dst_bits/src_bits\npad = (-x.shape[-2]) % (dst_bits // src_bits)\nx_p = jnp.pad(x, [(0,0),(0,pad)])\ny = mosaic.bitcast(x_p, jnp.float8_e4m3)","handlingStrategy":"validation","validationCode":"from jax import dtypes\ndef packs_evenly(x, ty):\n  return (x.shape[-2] * dtypes.itemsize_bits(x.dtype)) % dtypes.itemsize_bits(ty) == 0\nassert packs_evenly(x, jnp.float8_e4m3), \"2nd-minor dim cannot be packed evenly\"","typeGuard":null,"tryCatchPattern":null,"preventionTips":["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"],"tags":["jax","pallas","mosaic","bitcast","dtype","packing","shape-validation"],"backgroundTag":"dtype-bitcast-shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}