{"record":{"id":"d6f3299844bd403f","repo":"jax-ml/jax","slug":"conv-general-dilated-batch-group-count-must-divide","errorCode":null,"errorMessage":"conv_general_dilated batch_group_count must divide lhs batch dimension size, but {} does not divide {}.","messagePattern":"conv_general_dilated batch_group_count must divide lhs batch dimension size, but (.+?) does not divide (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/convolution.py","lineNumber":430,"sourceCode":"           \"feature_group_count must equal the rhs input feature dimension \"\n           \"size, but {} // {} != {}.\")\n    raise ValueError(msg.format(lhs_feature_count, feature_group_count,\n                                rhs.shape[dimension_numbers.rhs_spec[1]]))\n  if rhs.shape[dimension_numbers.rhs_spec[0]] % feature_group_count:\n    msg = (\"conv_general_dilated rhs output feature dimension size must be a \"\n           \"multiple of feature_group_count, but {} is not a multiple of {}.\")\n    raise ValueError(msg.format(rhs.shape[dimension_numbers.rhs_spec[0]],\n                                feature_group_count))\n\n  if not batch_group_count > 0:\n    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","sourceCodeStart":412,"sourceCodeEnd":448,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/convolution.py#L412-L448","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Set batch_group_count so it divides the lhs batch size exactly (e.g. 1, 2, or the batch size itself)","If you meant grouped filters over channels, use feature_group_count instead of batch_group_count","Drop or pad the batch to a multiple of batch_group_count before the convolution"],"exampleFix":"# before\ny = lax.conv_general_dilated(lhs, rhs, ..., batch_group_count=4)  # lhs batch = 6\n# after\ny = lax.conv_general_dilated(lhs, rhs, ..., batch_group_count=3)  # 6 % 3 == 0","handlingStrategy":"validation","validationCode":"assert lhs.shape[0] % batch_group_count == 0, f'batch {lhs.shape[0]} not divisible by {batch_group_count}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Compute batch_group_count from the actual batch size: lhs.shape[0] // n","Centralize grouped-conv config in one helper that validates divisibility"],"tags":["jax","convolution","shape-validation","grouped-conv"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}