jax-ml/jax · error · ValueError

conv_general_dilated batch_group_count must be a positive in

Error message

conv_general_dilated batch_group_count must be a positive integer, got {}.

What it means

batch_group_count in conv_general_dilated (used for vmap'd batched convs) must be a positive integer; the shape rule checks `batch_group_count > 0` and raises ValueError otherwise, mirroring the feature_group_count check.

Source

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

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

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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Set batch_group_count to a positive int (usually 1)
2. Prefer jax.vmap over conv for batching convs instead of manually setting batch_group_count
3. Audit any computed group parameters with max(1, n)

Example fix

// before
lax.conv_general_dilated(x, k, (1,1), 'SAME', batch_group_count=bgc)  # bgc == 0
// after
lax.conv_general_dilated(x, k, (1,1), 'SAME', batch_group_count=max(1, bgc))

assert batch_group_count > 0, batch_group_count

When it happens

Trigger: Calling lax.conv_general_dilated(..., batch_group_count=0) or with a negative value; typically the parameter is managed by jax.vmap's batch rules rather than set manually.

Common situations: Manually constructing batched convolutions instead of letting vmap handle batching; passing an uninitialized/default-zero groups variable into batch_group_count.

Related errors

• String padding is not implemented for transposed convolution
• padding argument to conv_general_dilated should be a string
• conv_general_dilated lhs and rhs must have the same number o
• conv_general_dilated feature_group_count must be a positive
• conv_general_dilated feature_group_count must divide lhs fea

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