jax-ml/jax · error · ValueError

At most one of batch_group_count and feature_group_count may

Error message

At most one of batch_group_count and feature_group_count may be > 1, got batch_group_count={} and feature_group_count={}

What it means

XLA convolutions support either batch grouping or feature grouping in a single op, not both simultaneously. The shape rule rejects any call where both batch_group_count>1 and feature_group_count>1.

Source

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

    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)
  return tuple(np.take(out_trans, np.argsort(out_perm)))


def _conv_general_dilated_sharding_rule(
    lhs: core.ShapedArray, rhs: core.ShapedArray, *, window_strides, padding,
    lhs_dilation, rhs_dilation, dimension_numbers, feature_group_count,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Set exactly one of the two counts to 1 (usually keep feature_group_count for grouped convs)
2. Express the combined effect manually, e.g. reshape/split the batch and run separate grouped convolutions
3. Vmap over the batch instead of using batch_group_count

Example fix

# before
out = lax.conv_general_dilated(lhs, rhs, ..., feature_group_count=8, batch_group_count=2)
# after
out = jax.vmap(lambda a, b: lax.conv_general_dilated(a, b, ..., feature_group_count=8))(lhs, rhs)

assert not (batch_group_count > 1 and feature_group_count > 1), 'only one group count may exceed 1'

When it happens

Trigger: Passing feature_group_count=4 and batch_group_count=2 (both >1) to jax.lax.conv_general_dilated.

Common situations: Combining code from two examples (one depthwise via feature groups, one batch-grouped) into one call; attempting fused grouped+batched convs not supported by XLA.

Related errors

• conv_general_dilated batch_group_count must divide lhs batch
• conv_general_dilated rhs output feature dimension size must
• convolution dimension_numbers list/tuple must be length 3, g
• Grouped convolutions are not supported on Pallas Mosaic TPU
• {name} wrapped function must be passed at least one argument

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