jax-ml/jax · error · NotImplementedError

Bitcast 1D ref with bitwidth change is not supported.

Error message

Bitcast 1D ref with bitwidth change is not supported.

What it means

jax._src.state.utils.bitcast raises NotImplementedError when bitcasting between dtypes of different bitwidths on a ref with fewer than 2 dimensions. Width-changing bitcasts on 1D refs can't be expressed because there's no second-minor dimension to absorb the size change.

Source

Thrown at jax/_src/state/utils.py:89

        c if is_ref else ref_get(c, ())
        for is_ref, c in zip(is_const_ref, all_consts)
    ]
    return core.eval_jaxpr(jaxpr, all_consts, *args0, *args1)

  hoisted_jaxpr, _ = pe.trace_to_jaxpr(
      _hoist, ft.flatten_args(*in_avals),
      jaxpr.debug_info.with_unknown_names())
  assert not hoisted_jaxpr.consts, "All consts should have been converted to refs"
  return hoisted_jaxpr


def bitcast(x, dtype: DTypeLike):
  x_bitwidth = dtypes.itemsize_bits(x.dtype)
  y_bitwidth = dtypes.itemsize_bits(dtype)
  shape = list(x.shape)
  if x_bitwidth != y_bitwidth:
    if len(shape) < 2:
      raise NotImplementedError(
          "Bitcast 1D ref with bitwidth change is not supported."
      )
    # Note: this is only valid on TPU.
    if shape[-2] * x_bitwidth % y_bitwidth != 0:
      raise ValueError(
          "Expected input and output shapes are the same after multiplying"
          " the second-minor dimension by the bitwidths."
      )
  shape[-2] = shape[-2] * x_bitwidth // y_bitwidth
  if x_bitwidth < y_bitwidth:
    ratio = y_bitwidth // x_bitwidth
    x = x.reshape(*x.shape[:-2], x.shape[-2] // ratio, ratio, -1).swapaxes(
        -1, -2
    )
  y = lax.bitcast_convert_type(x, dtype)
  if x_bitwidth > y_bitwidth:
    y = y.swapaxes(-1, -2).reshape(shape)
  return y

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Reshape the ref to 2D first (e.g. (1, n) or (n//ratio, ratio)), bitcast, then reshape back
Use same-bitwidth dtypes (e.g. float32 <-> int32)
Use .get() and jax.numpy reinterpretation on values instead of the ref

Example fix

# before
y = bitcast(ref_1d, jnp.int8)  # ref is f32
# after
ref2 = ref.reshape(1, ref.shape[0])
y = bitcast(ref2, jnp.int8).reshape(-1)

Defensive patterns

Strategy: fallback

Validate before calling

import jax.numpy as jnp
from jax._src import dtypes
if dtypes.itemsize_bits(x.dtype) != dtypes.itemsize_bits(dtype) and len(x.shape) < 2:
    x = x.reshape(1, x.shape[0]) if x.ndim == 1 else x

Try / catch

try:
    y = bitcast(x, dtype)
except NotImplementedError:
    y = bitcast(x.reshape(1, -1), dtype).reshape(-1)

Prevention

Reshape 1D refs to 2D before width-changing bitcasts
Prefer same-width dtype pairs for 1D buffers

When it happens

Trigger: bitcast(ref_1d, dtype) where itemsize_bits differ, e.g. bitcasting a 1D float32 ref to float16 or int8.

Common situations: Viewing 1D byte buffers as a different-width dtype (common in TPU packing/serialization code).

Related errors

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