jax-ml/jax · error · ValueError

conv_general_dilated lhs feature dimension size divided by f

Error message

conv_general_dilated lhs feature dimension size divided by feature_group_count must equal the rhs input feature dimension size, but {} // {} != {}.

What it means

After dividing the lhs feature dimension by feature_group_count, the result must equal the rhs (kernel) input-feature dimension — each group convolves a slice of size lhs_feature_count/groups against the kernel. If definitely_equal fails, ValueError reports lhs_feature // groups != rhs_in_features.

Source

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

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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Shape the kernel as (..., lhs_in_channels // feature_group_count, out_channels)
2. Cross-check with PyTorch semantics: torch weight is (out//g, in//g, H, W); JAX rhs is (H, W, in//g, out)
3. Set feature_group_count=1 if you didn't intend grouping

Example fix

// before
k = jax.random.normal(key, (3, 3, 8, 16))  # 8 == full in-channels
lax.conv_general_dilated(x_c8, k, (1,1), 'SAME', feature_group_count=2)
// after
k = jax.random.normal(key, (3, 3, 4, 16))  # 4 == 8 // 2
lax.conv_general_dilated(x_c8, k, (1,1), 'SAME', feature_group_count=2)

per_group = lhs_in_channels // feature_group_count
assert kernel.shape[rhs_in_feature_axis] == per_group, (kernel.shape, per_group)

When it happens

Trigger: Grouped conv where kernel input-feature dim doesn't match channels-per-group, e.g. lhs has 8 channels, groups=2 (4 per group) but kernel shaped (H,W,8,Cout) instead of (H,W,4,Cout).

Common situations: Reusing an ungrouped kernel shape after adding feature_group_count; PyTorch->JAX ports where weight shape interacts differently with groups (torch divides both in and out channels by groups).

Related errors

• conv_general_dilated rhs output feature dimension size must
• 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/c14c985cab53ccc1. Report an issue: GitHub.