{"record":{"id":"a23920097eab1576","repo":"jax-ml/jax","slug":"not-implemented-bitcast-1d","errorCode":null,"errorMessage":"Not implemented: bitcast 1D","messagePattern":"Not implemented: bitcast 1D","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/primitives.py","lineNumber":69,"sourceCode":"\nIntDeviceId = int | jax.Array\nMultiDimDeviceId = tuple[IntDeviceId, ...] | dict[str | tuple[str, ...], IntDeviceId]\nRef = state.AbstractRef | state.TransformedRef\n\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","sourceCodeStart":51,"sourceCodeEnd":87,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/primitives.py#L51-L87","documentation":"The Mosaic bitcast primitive requires the input to have at least 2 dimensions because it packs/unpacks elements along the second-to-minor dimension. A 1D array has no such dimension, so reinterpretation is rejected with this ValueError from the public bitcast wrapper.","triggerScenarios":"Calling jax.experimental.pallas.mosaic.primitives.bitcast (or _bitcast_batch_rule hitting it under vmap) with a 1D array, e.g. bitcast(f32_array_of_shape=(N,), jnp.int32).","commonSituations":"Reinterpreting raw TCM buffer words between int32 and float32 inside a Pallas kernel without reshaping to (N, 1) or a 2D block first; converting packed sub-byte weights and forgetting the required trailing dims.","solutions":["Reshape the input to at least 2D before bitcast, e.g. x.reshape(N, 1) or x[..., None]","Use jax.lax.bitcast_convert_type (or numpy view) outside the Pallas kernel for ordinary dtype reinterpretation","Make sure the second-to-minor dim size times src bitwidth is a multiple of the destination bitwidth after reshaping"],"exampleFix":"# before\ny = mosaic.bitcast(x, jnp.int32)  # x.shape == (1024,)\n\n# after\ny = mosaic.bitcast(x.reshape(1024, 1), jnp.int32)","handlingStrategy":"validation","validationCode":"def bitcast_safe(x, ty):\n  import jax\n  assert len(x.shape) >= 2, \"mosaic bitcast needs ndim >= 2; reshape first\"\n  return mosaic.bitcast(x, ty)\n# or simply pre-reshape:\nx2 = x.reshape(*x.shape, 1) if x.ndim < 2 else x","typeGuard":"def is_bitcastable_shape(x) -> bool:\n  return len(x.shape) >= 2","tryCatchPattern":null,"preventionTips":["Always reshape 1D buffers to 2D before mosaic bitcast","Use jax.lax.bitcast_convert_type outside Pallas kernels"],"tags":["jax","pallas","mosaic","bitcast","dtype","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"}