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.

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…
• Attempting to convert array of shape
• axis is out of range.
• bitcast_convert_type with different bitwidths not supported…
• Cannot signal on a non-()-shaped semaphore

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

Appendix: 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)