{"record":{"id":"e73291efc06846fd","repo":"jax-ml/jax","slug":"at-most-one-of-batch-group-count-and-feature-group","errorCode":null,"errorMessage":"At most one of batch_group_count and feature_group_count may be > 1, got batch_group_count={} and feature_group_count={}","messagePattern":"At most one of batch_group_count and feature_group_count may be > 1, got batch_group_count=(.+?) and feature_group_count=(.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/convolution.py","lineNumber":441,"sourceCode":"    msg = (\"conv_general_dilated batch_group_count \"\n           \"must be a positive integer, got {}.\")\n    raise ValueError(msg.format(batch_group_count))\n  lhs_batch_count = lhs.shape[dimension_numbers.lhs_spec[0]]\n  if batch_group_count > 1 and lhs_batch_count % batch_group_count != 0:\n    msg = (\"conv_general_dilated batch_group_count must divide lhs batch \"\n           \"dimension size, but {} does not divide {}.\")\n    raise ValueError(msg.format(batch_group_count, lhs_batch_count))\n\n  if rhs.shape[dimension_numbers.rhs_spec[0]] % batch_group_count:\n    msg = (\"conv_general_dilated rhs output feature dimension size must be a \"\n           \"multiple of batch_group_count, but {} is not a multiple of {}.\")\n    raise ValueError(msg.format(rhs.shape[dimension_numbers.rhs_spec[0]],\n                                batch_group_count))\n\n  if batch_group_count > 1 and feature_group_count > 1:\n    msg = (\"At most one of batch_group_count and feature_group_count may be > \"\n           \"1, got batch_group_count={} and feature_group_count={}\")\n    raise ValueError(msg.format(batch_group_count, feature_group_count))\n\n  if len(_conv_sdims(dimension_numbers.rhs_spec)) != len(window_strides):\n    msg = (\"conv_general_dilated window and window_strides must have \"\n           \"the same number of dimensions, but got {} and {}\")\n    raise ValueError(\n        msg.format(len(_conv_sdims(dimension_numbers.rhs_spec)), len(window_strides)))\n\n  lhs_perm, rhs_perm, out_perm = dimension_numbers\n  lhs_trans = lax._dilate_shape(np.take(lhs.shape, lhs_perm), lhs_dilation)\n  rhs_trans = lax._dilate_shape(np.take(rhs.shape, rhs_perm), rhs_dilation)\n  out_trans = conv_shape_tuple(lhs_trans, rhs_trans, window_strides, padding,\n                               batch_group_count)\n  return tuple(np.take(out_trans, np.argsort(out_perm)))\n\n\ndef _conv_general_dilated_sharding_rule(\n    lhs: core.ShapedArray, rhs: core.ShapedArray, *, window_strides, padding,\n    lhs_dilation, rhs_dilation, dimension_numbers, feature_group_count,","sourceCodeStart":423,"sourceCodeEnd":459,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/convolution.py#L423-L459","documentation":"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.","triggerScenarios":"Passing feature_group_count=4 and batch_group_count=2 (both >1) to jax.lax.conv_general_dilated.","commonSituations":"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.","solutions":["Set exactly one of the two counts to 1 (usually keep feature_group_count for grouped convs)","Express the combined effect manually, e.g. reshape/split the batch and run separate grouped convolutions","Vmap over the batch instead of using batch_group_count"],"exampleFix":"# before\nout = lax.conv_general_dilated(lhs, rhs, ..., feature_group_count=8, batch_group_count=2)\n# after\nout = jax.vmap(lambda a, b: lax.conv_general_dilated(a, b, ..., feature_group_count=8))(lhs, rhs)","handlingStrategy":"validation","validationCode":"assert not (batch_group_count > 1 and feature_group_count > 1), 'only one group count may exceed 1'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Default both counts to 1 and set only the one you need","Use vmap for batch handling when feature grouping is required"],"tags":["jax","convolution","grouped-conv","api-misuse"],"backgroundTag":"invalid-argument-combination","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}