jax-ml/jax · error · ValueError

conv_general_dilated rhs output feature dimension size must

Error message

conv_general_dilated rhs output feature dimension size must be a multiple of batch_group_count, but {} is not a multiple of {}.

What it means

For grouped-batch convolutions, the rhs (kernel) output-feature dimension is conceptually split across batch groups, so its size must be a multiple of batch_group_count. The shape rule enforces rhs.shape[dimension_numbers.rhs_spec[0]] % batch_group_count == 0.

Source

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

    msg = ("conv_general_dilated rhs output feature dimension size must be a "
           "multiple of feature_group_count, but {} is not a multiple of {}.")
    raise ValueError(msg.format(rhs.shape[dimension_numbers.rhs_spec[0]],
                                feature_group_count))

  if not batch_group_count > 0:
    msg = ("conv_general_dilated batch_group_count "
           "must be a positive integer, got {}.")
    raise ValueError(msg.format(batch_group_count))
  lhs_batch_count = lhs.shape[dimension_numbers.lhs_spec[0]]
  if batch_group_count > 1 and lhs_batch_count % batch_group_count != 0:
    msg = ("conv_general_dilated batch_group_count must divide lhs batch "
           "dimension size, but {} does not divide {}.")
    raise ValueError(msg.format(batch_group_count, lhs_batch_count))

  if rhs.shape[dimension_numbers.rhs_spec[0]] % batch_group_count:
    msg = ("conv_general_dilated rhs output feature dimension size must be a "
           "multiple of batch_group_count, but {} is not a multiple of {}.")
    raise ValueError(msg.format(rhs.shape[dimension_numbers.rhs_spec[0]],
                                batch_group_count))

  if batch_group_count > 1 and feature_group_count > 1:
    msg = ("At most one of batch_group_count and feature_group_count may be > "
           "1, got batch_group_count={} and feature_group_count={}")
    raise ValueError(msg.format(batch_group_count, feature_group_count))

  if len(_conv_sdims(dimension_numbers.rhs_spec)) != len(window_strides):
    msg = ("conv_general_dilated window and window_strides must have "
           "the same number of dimensions, but got {} and {}")
    raise ValueError(
        msg.format(len(_conv_sdims(dimension_numbers.rhs_spec)), len(window_strides)))

  lhs_perm, rhs_perm, out_perm = dimension_numbers
  lhs_trans = lax._dilate_shape(np.take(lhs.shape, lhs_perm), lhs_dilation)
  rhs_trans = lax._dilate_shape(np.take(rhs.shape, rhs_perm), rhs_dilation)
  out_trans = conv_shape_tuple(lhs_trans, rhs_trans, window_strides, padding,
                               batch_group_count)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make the rhs output-feature dimension a multiple of batch_group_count (e.g. set filters = filters rounded up)
2. Use feature_group_count for channel grouping instead of batch_group_count
3. Recheck the ConvDimensionNumbers spec — rhs_spec[0] may not be the axis you think is the feature dim

Example fix

# before
rhs = jnp.ones((3, 3, 10, 8))  # Cout=10
... = lax.conv_general_dilated(lhs, rhs, ..., batch_group_count=4)
# after
rhs = jnp.ones((3, 3, 12, 8))  # Cout=12 divisible by 4
... = lax.conv_general_dilated(lhs, rhs, ..., batch_group_count=4)

cout = rhs.shape[rhs_spec.index('O')] if isinstance(dn, str) else rhs.shape[-1]
assert cout % batch_group_count == 0

When it happens

Trigger: jax.lax.conv_general_dilated with batch_group_count>1 and a kernel whose output feature dim (e.g. weight shape (H,W,Cout,Cin) with Cout not divisible by batch_group_count).

Common situations: Porting flax equivariant/grouped layers with mismatched filter counts; typo where batch_group_count was set to the channel count instead of feature_group_count.

Related errors

• conv_general_dilated batch_group_count must divide lhs batch
• At most one of batch_group_count and feature_group_count may
• conv_general_dilated window and window_strides must have the
• Wrong number of explicit pads for convolution: expected {},
• Negative padding is larger than the size of the correspondin

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