{"record":{"id":"2321497afa6e5c5a","repo":"jax-ml/jax","slug":"mapped-axes-must-have-same-shape-got-in1-shape","errorCode":null,"errorMessage":"mapped axes must have same shape; got {in1.shape=} {in2.shape=} {axes=}","messagePattern":"mapped axes must have same shape; got (.+?) (.+?) (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/scipy/signal.py","lineNumber":114,"sourceCode":"\n    >>> with jax.numpy.printoptions(precision=3):\n    ...   print(jax.scipy.signal.fftconvolve(x, y, mode='valid'))\n    [6. 7. 6.]\n  \"\"\"\n  check_arraylike('fftconvolve', in1, in2)\n  in1, in2 = promote_dtypes_inexact(in1, in2)\n  if in1.ndim != in2.ndim:\n    raise ValueError(\"in1 and in2 should have the same dimensionality\")\n  if mode not in [\"same\", \"full\", \"valid\"]:\n    raise ValueError(\"mode must be one of ['same', 'full', 'valid']\")\n  _fftconvolve = partial(_fftconvolve_unbatched, mode=mode)\n  if axes is None:\n    return _fftconvolve(in1, in2)\n  axes = _ensure_index_tuple(axes)\n  axes = tuple(canonicalize_axis(ax, in1.ndim) for ax in axes)\n  mapped_axes = set(range(in1.ndim)) - set(axes)\n  if any(in1.shape[i] != in2.shape[i] for i in mapped_axes):\n    raise ValueError(f\"mapped axes must have same shape; got {in1.shape=} {in2.shape=} {axes=}\")\n  for ax in sorted(mapped_axes):\n    _fftconvolve = api.vmap(_fftconvolve, in_axes=ax, out_axes=ax)\n  return _fftconvolve(in1, in2)\n\ndef _fftconvolve_unbatched(in1: Array, in2: Array, mode: str) -> Array:\n  full_shape = tuple(s1 + s2 - 1 for s1, s2 in zip(in1.shape, in2.shape))\n\n  # TODO(jakevdp): potentially use next_fast_len to evaluate with a more efficient shape.\n  fft_shape = full_shape  # tuple(next_fast_len(s) for s in full_shape)\n\n  if mode == 'valid':\n    no_swap = all(s1 >= s2 for s1, s2 in zip(in1.shape, in2.shape))\n    swap = all(s1 <= s2 for s1, s2 in zip(in1.shape, in2.shape))\n    if not (no_swap or swap):\n      raise ValueError(\"For 'valid' mode, One input must be at least as \"\n                       \"large as the other in every dimension.\")\n    if swap:\n      in1, in2 = in2, in1","sourceCodeStart":96,"sourceCodeEnd":132,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/scipy/signal.py#L96-L132","documentation":"When fftconvolve is called with explicit axes, the remaining (mapped/batched) axes are vmap-ed over and must match in size between in1 and in2. Otherwise the batched convolution is ill-defined.","triggerScenarios":"Convolving (8, 100) and (16, 100) along axes=1 — batch dims 8 vs 16 differ.","commonSituations":"Batched signals where one side has a different batch size after slicing or padding.","solutions":["Make batch (non-axes) dims equal: slice, pad, or broadcast one input","Alternatively vmap manually if you actually want per-sample behavior","Double-check shapes after preprocessing pipelines"],"exampleFix":"# before\nsignal.fftconvolve(a, b, axes=1)  # a:(8,100), b:(16,100)\n# after\nb = b[:8]  # align batch dims\nsignal.fftconvolve(a, b, axes=1)","handlingStrategy":"validation","validationCode":"mapped = set(range(in1.ndim)) - set(axes)\nassert all(in1.shape[i] == in2.shape[i] for i in mapped)","typeGuard":"def batch_axes_match(in1, in2, axes) -> bool:\n    m = set(range(in1.ndim)) - set(axes)\n    return all(in1.shape[i] == in2.shape[i] for i in m)","tryCatchPattern":null,"preventionTips":["Print both shapes before batched convolutions in debug builds"],"tags":["jax","scipy","signal","vmap","shape-mismatch"],"backgroundTag":"batch-shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}