jax-ml/jax · error · ValueError

type of weights must match type of x. Got typeof(x)={core.ty

Error message

type of weights must match type of x. Got typeof(x)={core.typeof(x).str_short(True, True)} and typeof(weights)={core.typeof(weights).str_short(True, True)}

What it means

When weights are given to jnp.bincount, they must have the same shape as x and, if both are sharded across a mesh, compatible shardings. This ValueError reports a shape or sharding mismatch between x and weights.

Source

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

  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
    wts = core.typeof(weights).sharding
    if (np.shape(x) != np.shape(weights) or
        (not xts.mesh.empty and not wts.mesh.empty and xts != wts)):
      raise ValueError(
          "type of weights must match type of x. Got"
          f" typeof(x)={core.typeof(x).str_short(True, True)} and"
          f" typeof(weights)={core.typeof(weights).str_short(True, True)}")
  out_sharding = canonicalize_sharding(out_sharding, 'jnp.bincount')
  if out_sharding is not None and not is_replicated_or_unreduced(out_sharding):
    raise core.ShardingTypeError(
        "out_sharding passed to `jnp.bincount` can only be fully replicated"
        " or fully unreduced along all mesh axes")
  return array_creation.zeros(length, _dtype(weights)).at[clip(x, 0)].add(
      weights, mode='drop', out_sharding=out_sharding)


@overload
def broadcast_shapes(*shapes: Sequence[int]) -> tuple[int, ...]:
  ...

@overload
def broadcast_shapes(*shapes: Sequence[int | core.Tracer]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Match shapes exactly: weights must broadcast-free equal x.shape (e.g. weights of shape x.shape, or omit for uniform weights)
2. Check sharding: ensure both arrays are replicated or sharded identically on the same mesh before bincount
3. If weights are per-class, expand them: weights[class_ids][x] to get per-sample weights

Example fix

// before
jnp.bincount(pred_classes, weights=class_weights)  # class_weights: (num_classes,)
// after
jnp.bincount(pred_classes, weights=class_weights[pred_classes])

if weights is not None:
    assert np.shape(weights) == np.shape(x), 'weights shape must equal x shape'

When it happens

Trigger: jnp.bincount(x, weights=w) with w.shape != x.shape, or under jax.sharding mesh computation where x and weights carry different sharding types (different meshes or shard specs).

Common situations: Passing per-class weights of length num_classes instead of per-sample weights; multi-host/mesh pipelines where inputs were sharded differently before reaching bincount.

Related errors

• Length of sharding.spec ({len(out_s.spec)}) must be equal to
• Mismatch in result shapes. %s vs %s
• Axis must be specified when shapes of a and weights differ.
• Shape of weights must be consistent with shape of a along sp
• Weights shape must match 'a' shape when axis is None.

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