jax-ml/jax · error · ValueError

Not implemented: bitcast 1D

Error message

Not implemented: bitcast 1D

What it means

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.

Source

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

IntDeviceId = int | jax.Array
MultiDimDeviceId = tuple[IntDeviceId, ...] | dict[str | tuple[str, ...], IntDeviceId]
Ref = state.AbstractRef | state.TransformedRef


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)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reshape the input to at least 2D before bitcast, e.g. x.reshape(N, 1) or x[..., None]
2. Use jax.lax.bitcast_convert_type (or numpy view) outside the Pallas kernel for ordinary dtype reinterpretation
3. Make sure the second-to-minor dim size times src bitwidth is a multiple of the destination bitwidth after reshaping

Example fix

# before
y = mosaic.bitcast(x, jnp.int32)  # x.shape == (1024,)

# after
y = mosaic.bitcast(x.reshape(1024, 1), jnp.int32)

def bitcast_safe(x, ty):
  import jax
  assert len(x.shape) >= 2, "mosaic bitcast needs ndim >= 2; reshape first"
  return mosaic.bitcast(x, ty)
# or simply pre-reshape:
x2 = x.reshape(*x.shape, 1) if x.ndim < 2 else x

def is_bitcastable_shape(x) -> bool:
  return len(x.shape) >= 2

When it happens

Trigger: 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).

Common situations: 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.

Related errors

• Not implemented: the 2nd minor dim can not be perfectly pack
• Attempting to convert array of shape {operand.shape} from {o
• bitcast_convert_type with different bitwidths not supported
• {uninitialized_memory} + {str(dtype)}
• axis is out of range.

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