jax-ml/jax · error · ValueError

Wrong number of strides for spatial dimensions

Error message

Wrong number of strides for spatial dimensions

What it means

In the NumPy reference implementation, the number of window strides must equal the number of spatial dimensions (rhs rank minus 2). Passing one stride per non-spatial axis, or vice versa, triggers this ValueError.

Source

Thrown at jax/_src/lax_reference.py:457

                 in zip(out_shape, window_strides, filter_shape, in_shape)]
    if padding.upper() == 'SAME':
      return [
          (pad_size // 2, pad_size - pad_size // 2) for pad_size in pad_sizes
      ]
    else:
      return [
          (pad_size - pad_size // 2, pad_size // 2) for pad_size in pad_sizes
      ]
  else:
    return [(0, 0)] * len(in_shape)

def _conv_view(lhs, rhs_shape, window_strides, pads, pad_value):
  """Compute the view (and its axes) of a convolution or window reduction."""
  if (_min(lhs.ndim, len(rhs_shape)) < 2 or lhs.ndim != len(rhs_shape)
      or lhs.shape[1] != rhs_shape[1]):
    raise ValueError('Dimension mismatch')
  if len(window_strides) != len(rhs_shape) - 2:
    raise ValueError('Wrong number of strides for spatial dimensions')
  if len(pads) != len(rhs_shape) - 2:
    raise ValueError('Wrong number of pads for spatial dimensions')

  lhs = _pad(lhs, [(0, 0)] * 2 + list(pads), pad_value)
  in_shape = lhs.shape[2:]
  filter_shape = rhs_shape[2:]
  dim = len(filter_shape)  # number of 'spatial' dimensions in convolution

  out_strides = np.multiply(window_strides, lhs.strides[2:])
  view_strides = lhs.strides[:1] + tuple(out_strides) + lhs.strides[1:]

  out_shape = np.floor_divide(
      np.subtract(in_shape, filter_shape), window_strides) + 1
  view_shape = lhs.shape[:1] + tuple(out_shape) + rhs_shape[1:]

  view = np.lib.stride_tricks.as_strided(lhs, view_shape, view_strides)

  view_axes = list(range(view.ndim))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass strides only for spatial dimensions (drop batch/channel entries)
2. Check padding count matches too (len(pads) == spatial dims)

Example fix

# before
lax_reference.reduce_window(x, dims, (1,1,1,1), pads)  # 4 strides, 2 spatial
# after
lax_reference.reduce_window(x, dims, (1,1), pads)

assert len(window_strides) == len(rhs_shape) - 2

When it happens

Trigger: lax_reference._conv/reduce_window with len(window_strides) != len(rhs_shape) - 2, e.g. 4 strides for a 2-spatial-dim conv.

Common situations: Passing full-rank strides (including batch/channel) to the reference path which expects only spatial strides.

Related errors

• conv_general_dilated window and window_strides must have the
• Dimension mismatch
• Wrong number of pads for spatial dimensions
• conv_general_dilated batch_group_count must divide lhs batch
• conv_general_dilated rhs output feature dimension size must

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