{"record":{"id":"360a46390fab84fc","repo":"jax-ml/jax","slug":"type-of-weights-must-match-type-of-x-got-typeof-x","errorCode":null,"errorMessage":"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)}","messagePattern":"type of weights must match type of x\\. Got typeof\\(x\\)=(.+?) and typeof\\(weights\\)=(.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/lax_numpy.py","lineNumber":2985,"sourceCode":"  if length is None:\n    x_arr = core.concrete_or_error(\n        asarray, x,\n        \"The error occurred because of argument 'x' of jnp.bincount. \"\n        \"To avoid this error, pass a static `length` argument.\")\n    length = max(minlength, x_arr.size and int(max(0, x_arr.max())) + 1)\n  else:\n    length = core.concrete_dim_or_error(\n        length,\n        \"The error occurred because of argument 'length' of jnp.bincount.\")\n\n  if weights is None:\n    weights = np.array(1, dtype=dtypes.int_)\n  else:\n    xts = core.typeof(x).sharding\n    wts = core.typeof(weights).sharding\n    if (np.shape(x) != np.shape(weights) or\n        (not xts.mesh.empty and not wts.mesh.empty and xts != wts)):\n      raise ValueError(\n          \"type of weights must match type of x. Got\"\n          f\" typeof(x)={core.typeof(x).str_short(True, True)} and\"\n          f\" typeof(weights)={core.typeof(weights).str_short(True, True)}\")\n  out_sharding = canonicalize_sharding(out_sharding, 'jnp.bincount')\n  if out_sharding is not None and not is_replicated_or_unreduced(out_sharding):\n    raise core.ShardingTypeError(\n        \"out_sharding passed to `jnp.bincount` can only be fully replicated\"\n        \" or fully unreduced along all mesh axes\")\n  return array_creation.zeros(length, _dtype(weights)).at[clip(x, 0)].add(\n      weights, mode='drop', out_sharding=out_sharding)\n\n\n@overload\ndef broadcast_shapes(*shapes: Sequence[int]) -> tuple[int, ...]:\n  ...\n\n@overload\ndef broadcast_shapes(*shapes: Sequence[int | core.Tracer]","sourceCodeStart":2967,"sourceCodeEnd":3003,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/lax_numpy.py#L2967-L3003","documentation":"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.","triggerScenarios":"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).","commonSituations":"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.","solutions":["Match shapes exactly: weights must broadcast-free equal x.shape (e.g. weights of shape x.shape, or omit for uniform weights)","Check sharding: ensure both arrays are replicated or sharded identically on the same mesh before bincount","If weights are per-class, expand them: weights[class_ids][x] to get per-sample weights"],"exampleFix":"// before\njnp.bincount(pred_classes, weights=class_weights)  # class_weights: (num_classes,)\n// after\njnp.bincount(pred_classes, weights=class_weights[pred_classes])","handlingStrategy":"validation","validationCode":"if weights is not None:\n    assert np.shape(weights) == np.shape(x), 'weights shape must equal x shape'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Per-class weights must be gathered by sample: class_weights[x]","Under mesh/sharding, ensure x and weights have identical shardings before bincount"],"tags":["jax","bincount","weights","sharding","shape-mismatch"],"backgroundTag":"argument-shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}