jax-ml/jax · error · ValueError

conv_general_dilated lhs and rhs must have the same number o

Error message

conv_general_dilated lhs and rhs must have the same number of dimensions, but got {} and {}.

What it means

The shape rule for conv_general_dilated requires lhs (input) and rhs (kernel) to have identical rank, since a single ConvDimensionNumbers spec maps dimensions of both. If len(lhs.shape) != len(rhs.shape) it raises ValueError showing both shapes.

Source

Thrown at jax/_src/lax/convolution.py:399

      pads = padding
  if transpose_kernel:
    # flip spatial dims and swap input / output channel axes
    rhs = _flip_axes(rhs, np.array(dn.rhs_spec)[2:])
    rhs = rhs.swapaxes(dn.rhs_spec[0], dn.rhs_spec[1])
  return conv_general_dilated(lhs, rhs, one, pads, strides, rhs_dilation, dn,
                              precision=precision,
                              preferred_element_type=preferred_element_type)


def _conv_general_dilated_shape_rule(
    lhs: core.ShapedArray, rhs: core.ShapedArray, *, window_strides, padding,
    lhs_dilation, rhs_dilation, dimension_numbers, feature_group_count,
    batch_group_count, **unused_kwargs) -> tuple[int, ...]:
  assert type(dimension_numbers) is ConvDimensionNumbers
  if len(lhs.shape) != len(rhs.shape):
    msg = ("conv_general_dilated lhs and rhs must have the same number of "
           "dimensions, but got {} and {}.")
    raise ValueError(msg.format(lhs.shape, rhs.shape))
  if not feature_group_count > 0:
    msg = ("conv_general_dilated feature_group_count "
           "must be a positive integer, got {}.")
    raise ValueError(msg.format(feature_group_count))
  lhs_feature_count = lhs.shape[dimension_numbers.lhs_spec[1]]
  quot, rem = divmod(lhs_feature_count, feature_group_count)
  if rem:
    msg = ("conv_general_dilated feature_group_count must divide lhs feature "
           "dimension size, but {} does not divide {}.")
    raise ValueError(msg.format(feature_group_count, lhs_feature_count))
  if not core.definitely_equal(quot, rhs.shape[dimension_numbers.rhs_spec[1]]):
    msg = ("conv_general_dilated lhs feature dimension size divided by "
           "feature_group_count must equal the rhs input feature dimension "
           "size, but {} // {} != {}.")
    raise ValueError(msg.format(lhs_feature_count, feature_group_count,
                                rhs.shape[dimension_numbers.rhs_spec[1]]))
  if rhs.shape[dimension_numbers.rhs_spec[0]] % feature_group_count:
    msg = ("conv_general_dilated rhs output feature dimension size must be a "

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Add/remove leading dimensions to match ranks: x[None] for the batch dim, or kernel as (H,W,Cin,Cout)
2. Check both .ndim before calling: assert lhs.ndim == rhs.ndim
3. Use convenience wrappers (jnp.convolution) that handle common ranks

Example fix

// before
out = lax.conv(x_hwc, kernel_hwio, (1,1), 'SAME')  # ranks 3 vs 4
// after
out = lax.conv(x_nhwc, kernel_hwio, (1,1), 'SAME')  # x[None] added

assert lhs.ndim == rhs.ndim, (lhs.shape, rhs.shape)

When it happens

Trigger: Convolving a 4-D input (NHWC) with a 3-D kernel (HWI) or similar rank mismatch in lax.conv/conv_general_dilated; often via the Python-level wrapper passing unbatched kernels.

Common situations: Forgetting to add the batch dimension (passing HWC input against HWIO kernel); sharing weights across a batch but reshaping the kernel incorrectly; vmap or tree_map stripping a leading dim from only one operand.

Related errors

• String padding is not implemented for transposed convolution
• padding argument to conv_general_dilated should be a string
• conv_general_dilated feature_group_count must be a positive
• conv_general_dilated feature_group_count must divide lhs fea
• conv_general_dilated lhs feature dimension size divided by f

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