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 feature_group_count, but {} is not a multiple of {}.

What it means

In grouped convolutions the rhs (kernel) output-feature dimension must be a multiple of feature_group_count so outputs can be split evenly across groups. The shape rule checks rhs out-features % feature_group_count and raises ValueError naming both values.

Source

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

    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 "
           "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))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make out_channels divisible by feature_group_count (round up/down to nearest multiple)
2. For depthwise conv set out_channels = in_channels * multiplier so divisibility holds
3. Validate at model-build time: assert cout % groups == 0

Example fix

// before
lax.conv_general_dilated(x, k_cout6, (1,1), 'SAME', feature_group_count=4)
// after
k_cout8 = ...  # 8 % 4 == 0
lax.conv_general_dilated(x, k_cout8, (1,1), 'SAME', feature_group_count=4)

assert kernel.shape[rhs_out_feature_axis] % feature_group_count == 0

When it happens

Trigger: Grouped conv with kernel output channels not divisible by groups, e.g. groups=4 with a kernel producing 6 output channels.

Common situations: Choosing output channels via a width multiplier that breaks divisibility by groups; ports from PyTorch where out_channels is already per-multiple of groups but rounding differs.

Related errors

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

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