jax-ml/jax · error · ValueError

conv_general_dilated batch_group_count must divide lhs batch

Error message

conv_general_dilated batch_group_count must divide lhs batch dimension size, but {} does not divide {}.

What it means

In lax.conv_general_dilated, batch_group_count splits the lhs batch dimension into groups (used for grouped/depthwise-style batched convolutions). The shape rule requires the lhs batch dimension size to be evenly divisible by batch_group_count. The message reports the offending batch_group_count and the actual lhs batch size.

Source

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

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

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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Set batch_group_count so it divides the lhs batch size exactly (e.g. 1, 2, or the batch size itself)
2. If you meant grouped filters over channels, use feature_group_count instead of batch_group_count
3. Drop or pad the batch to a multiple of batch_group_count before the convolution

Example fix

# before
y = lax.conv_general_dilated(lhs, rhs, ..., batch_group_count=4)  # lhs batch = 6
# after
y = lax.conv_general_dilated(lhs, rhs, ..., batch_group_count=3)  # 6 % 3 == 0

assert lhs.shape[0] % batch_group_count == 0, f'batch {lhs.shape[0]} not divisible by {batch_group_count}'

When it happens

Trigger: Calling jax.lax.conv_general_dilated (or lax.conv with batch_group_count) with batch_group_count>1 where lhs.shape[dimension_numbers.lhs_spec[0]] % batch_group_count != 0, e.g. batch of 5 with batch_group_count=2.

Common situations: Migrating from TensorFlow grouped convolutions and passing a group count larger than the batch; using batch_group_count where feature_group_count was intended; dynamic batch sizes (last partial batch) that are not multiples of the group count.

Related errors

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