jax-ml/jax · error · ValueError

conv_general_dilated feature_group_count must divide lhs fea

Error message

conv_general_dilated feature_group_count must divide lhs feature dimension size, but {} does not divide {}.

What it means

For grouped convolutions, the lhs (input) feature dimension must be divisible by feature_group_count so each group gets an equal slice of input channels. The shape rule divmods lhs feature count by the group count and raises ValueError when a remainder exists.

Source

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

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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Choose feature_group_count that divides the input channel count (common divisors: 1, 2, in_channels//k)
Adjust the input channels to a multiple of the group count
For depthwise: feature_group_count = in_channels and rhs input-feature dim = 1

Example fix

// before
lax.conv_general_dilated(x_c6, k, (1,1), 'SAME', feature_group_count=4)
// after
lax.conv_general_dilated(x_c6, k, (1,1), 'SAME', feature_group_count=3)  # 6 % 3 == 0

Defensive patterns

Strategy: validation

Validate before calling

in_ch = lhs.shape[lhs_spec_feature]
assert in_ch % feature_group_count == 0, (in_ch, feature_group_count)

Prevention

Choose groups among divisors of the input channel count
Re-validate groups whenever changing network width

When it happens

Trigger: lax.conv_general_dilated with feature_group_count=4 on an input with 6 input channels (6 % 4 != 0); groups not dividing the number of input channels.

Common situations: Porting PyTorch grouped convs with a wrong groups value; changing the channel width of a network (e.g. width multiplier) without updating groups; depthwise conv with feature_group_count != number of input channels.

Related errors

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