jax-ml/jax · error · ValueError

Dimension mismatch

Error message

Dimension mismatch

What it means

In jax.lax_reference (the pure-NumPy eager backend), _conv_view validates convolution/window-reduction inputs: lhs must be at least 2-D, match rhs rank, and lhs.shape[1] must equal rhs.shape[1] (channel dims). A violation raises 'Dimension mismatch'.

Source

Thrown at jax/_src/lax_reference.py:455

    pad_sizes = [_max((out_size - 1) * stride + filter_size - in_size, 0)
                 for out_size, stride, filter_size, in_size
                 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 on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reshape operands to expected NHWC/OIHW layout with matching channel dimension
2. Verify lhs.ndim == len(rhs_shape) >= 2 and lhs.shape[1] == rhs_shape[1] before calling
3. Run under jit where the XLA path gives detailed shape errors

assert lhs.ndim >= 2 and lhs.ndim == len(rhs_shape) and lhs.shape[1] == rhs_shape[1]

When it happens

Trigger: Calling lax_reference.reduce_window or _conv with a 1-D input, mismatched ranks, or a rhs whose channel dim differs from the lhs channel dim (e.g. NumPy conv path used with wrong filter layout).

Common situations: Using jax operations outside jit/tracing (lax_reference fallback) with data shaped for another framework (e.g. filters in HWIO vs OIHW).

Related errors

• Wrong number of strides for spatial dimensions
• 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
• conv_general_dilated window and window_strides must have the

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