{"record":{"id":"19926a24faaa8950","repo":"jax-ml/jax","slug":"weights-shape-weights-shape-must-match-reduction","errorCode":null,"errorMessage":"Weights shape {weights.shape} must match reduction axes {tuple(a.shape[ax] for ax in ax_tuple)}","messagePattern":"Weights shape (.+?) must match reduction axes (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/reductions.py","lineNumber":2530,"sourceCode":"    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):\n    keepdim = list(a.shape)\n    nd = a.ndim\n    axis = tuple(canonicalize_axis(ax, nd) for ax in axis)","sourceCodeStart":2512,"sourceCodeEnd":2548,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/reductions.py#L2512-L2548","documentation":"When an explicit axis is given to a weighted quantile, weights must have shape equal to the reduction axes of a (e.g. a.shape[axis] for a single axis). Otherwise JAX raises ValueError showing both shapes.","triggerScenarios":"Calling jnp.quantile(a, q, axis=0, weights=w, method='inverted_cdf') where w.shape != (a.shape[0],), e.g. full a-shaped weights with axis set, or transposed weights.","commonSituations":"Switching a working axis=None call to a per-axis reduction without reshaping weights; transposition bugs where weights align to the wrong dimension.","solutions":["Reshape weights to the reduction axes: w = w.reshape(a.shape[axis]) for a single axis","Or keep weights full-shaped and drop axis (axis=None) so they match a.shape","Double-check axis orientation vs weight layout with a shape assertion before calling"],"exampleFix":"// before\njnp.quantile(a, q, axis=0, weights=w, method='inverted_cdf')  # w has shape of full a\n// after\njnp.quantile(a, q, axis=0, weights=w.reshape(a.shape[0]), method='inverted_cdf')","handlingStrategy":"validation","validationCode":"import jax.numpy as jnp\n\nax = jnp.canonicalize_axis(axis, a.ndim)\nif weights.shape != (a.shape[ax],):\n    weights = weights.reshape(a.shape[ax])\njnp.quantile(a, q, axis=axis, weights=weights, method='inverted_cdf')","typeGuard":"def weights_match_reduction_axes(w, a, axis) -> bool:\n    ax = jnp.canonicalize_axis(axis, a.ndim)\n    return w.shape == (a.shape[ax],)","tryCatchPattern":null,"preventionTips":["For axis-wise weighted quantiles, keep weights 1-d over the reduction axis","Add shape assertions before stats calls on batched data","Beware transposed layouts when porting NumPy code"],"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"}