jax-ml/jax · error · ValueError

Cannot bitcast a ()-shaped array to a dtype with a different

Error message

Cannot bitcast a ()-shaped array to a dtype with a different bitwidth: {old_bitwidth=} vs {new_bitwidth=}

What it means

Bitcasting a scalar (ndim 0) array to a dtype with a different bitwidth is impossible because there is no trailing dimension to redistribute bits across, so the abstract eval raises.

Source

Thrown at jax/_src/pallas/mosaic/sc_primitives.py:442

    raise ValueError("Indices must not be empty")
  ref, transforms = state_primitives.get_ref_and_transforms(
      ref, None, "addupdate_scatter"
  )
  flat_args, tree = jax.tree.flatten((ref, transforms, indices, x, mask))
  _ = scatter_p.bind(*flat_args, tree=tree, add=True)


bitcast_p = jax_core.Primitive("bitcast")


@bitcast_p.def_abstract_eval
def _bitcast_abstract_eval(x, dtype):
  old_bitwidth = dtypes.itemsize_bits(x.dtype)
  new_bitwidth = dtypes.itemsize_bits(dtype)
  if old_bitwidth == new_bitwidth:
    return jax_core.ShapedArray(x.shape, dtype)
  if x.ndim == 0:
    raise ValueError(
        "Cannot bitcast a ()-shaped array to a dtype with a different bitwidth:"
        f" {old_bitwidth=} vs {new_bitwidth=}"
    )
  new_last_dim, rem = divmod(x.shape[-1] * old_bitwidth, new_bitwidth)
  if rem:
    raise ValueError(
        f"Cannot bitcast from {x.dtype} ({old_bitwidth} bits) to"
        f" {dtype} ({new_bitwidth} bits), because {x.shape[-1]=} *"
        f" {old_bitwidth} is not divisible by {new_bitwidth}"
    )
  return jax_core.ShapedArray((*x.shape[:-1], new_last_dim), dtype)


@sc_lowering.register_lowering_rule(bitcast_p)
def _bitcast_lowering_rule(ctx: sc_lowering.LoweringRuleContext, x, *, dtype):
  del dtype  # Unused.
  [out_aval] = ctx.avals_out
  return vector.bitcast(ctx.aval_to_ir_type(out_aval), x)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Reshape the scalar to shape (1,) before bitcasting (and reshape back after)
  2. Use lax.convert_element_type if you want a value conversion, not a bit reinterpretation
  3. Match bitwidths (e.g. f32<->u32) for scalar bitcasts

Example fix

// before
y = bitcast(jnp.float32(1.0), jnp.int16)  # scalar, bitwidth differs

// after
y = bitcast(jnp.float32(1.0).reshape(1), jnp.int16).reshape(())
Defensive patterns

Strategy: validation

Validate before calling

if x.ndim == 0 and dtypes.itemsize_bits(x.dtype) != dtypes.itemsize_bits(dtype):
    x = x.reshape(1)

Type guard

def bitcast_ok(x, dtype) -> bool:
    ob, nb = dtypes.itemsize_bits(x.dtype), dtypes.itemsize_bits(dtype)
    return ob == nb or (x.ndim > 0 and (x.shape[-1] * ob) % nb == 0)

Prevention

When it happens

Trigger: Calling the SC bitcast primitive on a 0-d array, e.g. bitcast(jnp.float32(1.0), jnp.int16) or to a wider dtype like float32->bfloat16 pairs on a scalar.

Common situations: Applying vectorized bitcast logic to scalars; converting constants/registers inside a kernel where a scalar slipped through.

Related errors


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