jax-ml/jax · error · ValueError

only 1-dimensional input supported.

Error message

only 1-dimensional input supported.

What it means

jnp.bincount only accepts 1-dimensional arrays; bincount is defined as a per-value histogram over a flat vector, and multi-dimensional input is rejected.

Source

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

    >>> jit_bincount = jax.jit(jnp.bincount, static_argnames=['length'])
    >>> jit_bincount(x, length=5)
    Array([0, 2, 1, 3, 0], dtype=int32)

    Any negative numbers are clipped to the first bin, and numbers beyond the
    specified ``length`` are dropped:

    >>> x = jnp.array([-1, -1, 1, 3, 10])
    >>> jnp.bincount(x, length=5)
    Array([2, 1, 0, 1, 0], dtype=int32)
  """
  x = util.ensure_arraylike("bincount", x)
  if x.dtype == bool:
    x = lax.convert_element_type(x, 'int32')
  if not issubdtype(x.dtype, np.integer):
    raise TypeError(f"x argument to bincount must have an integer type; got {x.dtype}")
  if np.ndim(x) != 1:
    raise ValueError("only 1-dimensional input supported.")
  minlength = core.concrete_or_error(
      operator.index, minlength,
      "The error occurred because of argument 'minlength' of jnp.bincount.")
  if length is None:
    x_arr = core.concrete_or_error(
        asarray, x,
        "The error occurred because of argument 'x' of jnp.bincount. "
        "To avoid this error, pass a static `length` argument.")
    length = max(minlength, x_arr.size and int(max(0, x_arr.max())) + 1)
  else:
    length = core.concrete_dim_or_error(
        length,
        "The error occurred because of argument 'length' of jnp.bincount.")

  if weights is None:
    weights = np.array(1, dtype=dtypes.int_)
  else:
    xts = core.typeof(x).sharding

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Flatten first: jnp.bincount(labels.ravel())
2. If you need per-row counts, vmap over rows: jax.vmap(lambda r: jnp.bincount(r, length=k))(labels)
3. Consider jnp.histogram-style or scatter-based counting for multi-dim use

Example fix

// before
jnp.bincount(labels)  # labels.shape == (B, N)
// after
jnp.bincount(labels.ravel())

if x.ndim != 1:
    x = x.ravel()

def is_1d(a) -> bool:
    return a.ndim == 1

When it happens

Trigger: Calling jnp.bincount on a 2D array, e.g. jnp.bincount(jnp.array([[1,2],[3,4]])), or on labels of shape (batch, seq_len).

Common situations: Batched label counting in ML pipelines where labels carry a batch dimension; forgetting to ravel after reductions.

Related errors

• Unsupported ndim: {x.ndim}
• Unsupported shape: {x.shape}
• scan got `length` argument of {} which disagrees with leadin
• conv_general_dilated batch_group_count must divide lhs batch
• conv_general_dilated rhs output feature dimension size must

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