jax-ml/jax · error · TypeError

Expected an input array of unsigned byte data type

Error message

Expected an input array of unsigned byte data type

What it means

Raised by jnp.unpackbits when the input dtype is not uint8 — unpackbits expands each byte into 8 bits, so only unsigned 8-bit input is meaningful.

Source

Thrown at jax/_src/numpy/lax_numpy.py:8789

           [ 49]], dtype=uint8)

    The ``count`` keyword lets ``unpackbits`` serve as an inverse of ``packbits``
    in cases where not all bits are present:

    >>> bits = jnp.array([1, 1, 0, 1, 1, 0, 1, 1, 0, 1, 1])  # 11 bits
    >>> vals = jnp.packbits(bits)
    >>> vals
    Array([219,  96], dtype=uint8)
    >>> jnp.unpackbits(vals)  # 16 zero-padded bits
    Array([1, 1, 0, 1, 1, 0, 1, 1, 0, 1, 1, 0, 0, 0, 0, 0], dtype=uint8)
    >>> jnp.unpackbits(vals, count=11)  # specify 11 output bits
    Array([1, 1, 0, 1, 1, 0, 1, 1, 0, 1, 1], dtype=uint8)
    >>> jnp.unpackbits(vals, count=-5)  # specify 5 bits to be trimmed
    Array([1, 1, 0, 1, 1, 0, 1, 1, 0, 1, 1], dtype=uint8)
  """
  arr = util.ensure_arraylike("unpackbits", a)
  if arr.dtype != np.uint8:
    raise TypeError("Expected an input array of unsigned byte data type")
  if bitorder not in ['little', 'big']:
    raise ValueError("'order' must be either 'little' or 'big'")
  bits = asarray(1) << arange(8, dtype='uint8')
  if bitorder == 'big':
    bits = bits[::-1]
  if axis is None:
    arr = ravel(arr)
    axis = 0
  arr = swapaxes(arr, axis, -1)
  unpacked = ((arr[..., None] & expand_dims(bits, tuple(range(arr.ndim)))) > 0).astype('uint8')
  unpacked = unpacked.reshape(unpacked.shape[:-2] + (-1,))
  if count is not None:
    if count > unpacked.shape[-1]:
      unpacked = pad(unpacked, [(0, 0)] * (unpacked.ndim - 1) + [(0, count - unpacked.shape[-1])])
    else:
      unpacked = unpacked[..., :count]
  return swapaxes(unpacked, axis, -1)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast explicitly: jnp.unpackbits(a.astype(np.uint8))
2. Avoid promotions upstream (use lax operations that preserve dtype, or ufuncs with dtype=uint8)
3. Confirm values fit in [0,255] before casting to avoid silent truncation

Example fix

// before
jnp.unpackbits(flags + 0)  # promoted to int32
// after
jnp.unpackbits((flags + 0).astype(jnp.uint8))

arr = jnp.asarray(a)
if arr.dtype != jnp.uint8: arr = arr.astype(jnp.uint8)

def is_uint8(a):
    return jnp.asarray(a).dtype == jnp.uint8

When it happens

Trigger: jnp.unpackbits(int32_array), float arrays, or int8/uint16 data; results of arithmetic that promoted uint8 to int32.

Common situations: JAX type promotion: (uint8_array + 0) or comparisons yielding wider dtypes; data loaded as int64 indices then passed to unpackbits; porting NumPy code that relied on implicit casting.

Related errors

• {} does not accept dtype {}. Accepted dtypes are subtypes of
• {name} does not accept dtype {dtype_to_string(aval.dtype)}.
• {} does not accept dtype {} at position {}. Accepted dtypes
• Input type is incompatible with `preferred_element_type`. Th
• `preferred_element_type` must have the same signedness as th

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