{"record":{"id":"a282ef1f01f09763","repo":"jax-ml/jax","slug":"weights-shape-must-match-a-shape-when-axis-is-no","errorCode":null,"errorMessage":"Weights shape must match 'a' shape when axis is None.","messagePattern":"Weights shape must match 'a' shape when axis is None\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/reductions.py","lineNumber":2527,"sourceCode":"  if overwrite_input or out is not None:\n    msg = (\"jax.numpy.nanquantile does not support overwrite_input=True or \"\n           \"out != None\")\n    raise ValueError(msg)\n  return _quantile(a, q, axis, method, keepdims, True, weights)\n\ndef _quantile(a: Array, q: Array, axis: int | tuple[int, ...] | None,\n              method: str, keepdims: bool, squash_nans: bool, weights: Array | None = None) -> Array:\n  if method not in [\"linear\", \"lower\", \"higher\", \"midpoint\", \"nearest\", \"inverted_cdf\"]:\n    raise ValueError(\"method can only be 'linear', 'lower', 'higher', 'midpoint', 'nearest' or 'inverted_cdf'\")\n  if weights is not None:\n    if dtypes.issubdtype(weights.dtype, np.complexfloating):\n      raise ValueError(\"Weights cannot be complex types.\")\n    if method != \"inverted_cdf\":\n      raise NotImplementedError(f\"{method} doesn't support weights. Only method 'inverted_cdf' supports weights.\")\n    a, weights = promote_dtypes_inexact(a, weights)\n    if weights.shape != a.shape:\n      if axis is None:\n        raise ValueError(\"Weights shape must match 'a' shape when axis is None.\")\n      ax_tuple = canonicalize_axis_tuple(axis, a.ndim)\n      if weights.shape != tuple(a.shape[ax] for ax in ax_tuple):\n        raise ValueError(f\"Weights shape {weights.shape} must match reduction axes \"\n                          f\"{tuple(a.shape[ax] for ax in ax_tuple)}\")\n      weights = lax.broadcast_in_dim(weights, a.shape, broadcast_dimensions=ax_tuple)\n  else:\n    a, = promote_dtypes_inexact(a)\n  keepdim = []\n  if dtypes.issubdtype(a.dtype, np.complexfloating):\n    raise ValueError(\"quantile does not support complex input, as the operation is poorly defined.\")\n  if axis is None:\n    if keepdims:\n      keepdim = [1] * a.ndim\n    a = a.ravel()\n    if weights is not None:\n      weights = weights.ravel()\n    axis = 0\n  elif isinstance(axis, tuple):","sourceCodeStart":2509,"sourceCodeEnd":2545,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/reductions.py#L2509-L2545","documentation":"When axis=None, weighted quantile requires weights with exactly the same shape as a (the reduction is over the whole array). A shape mismatch raises ValueError.","triggerScenarios":"Calling jnp.quantile(a, q, axis=None, weights=w, method='inverted_cdf') where w.shape != a.shape, e.g. flat weights against a 2-d array.","commonSituations":"Passing per-feature weight vectors to a whole-array quantile; reshaping a for a batched pipeline while keeping old 1-d weights.","solutions":["Broadcast weights to a.shape first: w = jnp.broadcast_to(w, a.shape)","Pass an explicit axis and supply weights matching just the reduction axes","Reshape weights: w.reshape(a.shape) when sizes match element-wise"],"exampleFix":"// before\njnp.quantile(a, q, weights=w, method='inverted_cdf')  # w.shape != a.shape, axis=None\n// after\njnp.quantile(a, q, weights=jnp.broadcast_to(w, a.shape), method='inverted_cdf')","handlingStrategy":"validation","validationCode":"import jax.numpy as jnp\n\nif axis is None and weights is not None and weights.shape != a.shape:\n    weights = jnp.broadcast_to(weights, a.shape)\njnp.quantile(a, q, axis=axis, weights=weights, method='inverted_cdf')","typeGuard":"def weights_match_full_shape(w, a) -> bool:\n    return w.shape == a.shape","tryCatchPattern":"try:\n    jnp.quantile(a, q, weights=w, method='inverted_cdf')\nexcept ValueError as e:\n    if 'Weights shape' in str(e):\n        w = jnp.broadcast_to(w, a.shape)\n        q_val = jnp.quantile(a, q, weights=w, method='inverted_cdf')\n    else:\n        raise","preventionTips":["Assert weights.shape == a.shape when axis is None","Centralize a weighted-quantile helper that normalizes weight shape"],"tags":["jax","numpy","quantile","weights","shape-mismatch"],"backgroundTag":"weight-shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}