{"record":{"id":"521896a5d63a5ef7","repo":"jax-ml/jax","slug":"conv-general-dilated-rhs-output-feature-dimension-521896","errorCode":null,"errorMessage":"conv_general_dilated rhs output feature dimension size must be a multiple of batch_group_count, but {} is not a multiple of {}.","messagePattern":"conv_general_dilated rhs output feature dimension size must be a multiple of batch_group_count, but (.+?) is not a multiple of (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/convolution.py","lineNumber":435,"sourceCode":"    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\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)","sourceCodeStart":417,"sourceCodeEnd":453,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/convolution.py#L417-L453","documentation":"For grouped-batch convolutions, the rhs (kernel) output-feature dimension is conceptually split across batch groups, so its size must be a multiple of batch_group_count. The shape rule enforces rhs.shape[dimension_numbers.rhs_spec[0]] % batch_group_count == 0.","triggerScenarios":"jax.lax.conv_general_dilated with batch_group_count>1 and a kernel whose output feature dim (e.g. weight shape (H,W,Cout,Cin) with Cout not divisible by batch_group_count).","commonSituations":"Porting flax equivariant/grouped layers with mismatched filter counts; typo where batch_group_count was set to the channel count instead of feature_group_count.","solutions":["Make the rhs output-feature dimension a multiple of batch_group_count (e.g. set filters = filters rounded up)","Use feature_group_count for channel grouping instead of batch_group_count","Recheck the ConvDimensionNumbers spec — rhs_spec[0] may not be the axis you think is the feature dim"],"exampleFix":"# before\nrhs = jnp.ones((3, 3, 10, 8))  # Cout=10\n... = lax.conv_general_dilated(lhs, rhs, ..., batch_group_count=4)\n# after\nrhs = jnp.ones((3, 3, 12, 8))  # Cout=12 divisible by 4\n... = lax.conv_general_dilated(lhs, rhs, ..., batch_group_count=4)","handlingStrategy":"validation","validationCode":"cout = rhs.shape[rhs_spec.index('O')] if isinstance(dn, str) else rhs.shape[-1]\nassert cout % batch_group_count == 0","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Round filter counts to multiples of the group count when constructing kernels","Prefer feature_group_count for channel grouping"],"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"}