{"record":{"id":"c14c985cab53ccc1","repo":"jax-ml/jax","slug":"conv-general-dilated-lhs-feature-dimension-size-di","errorCode":null,"errorMessage":"conv_general_dilated lhs feature dimension size divided by feature_group_count must equal the rhs input feature dimension size, but {} // {} != {}.","messagePattern":"conv_general_dilated lhs feature dimension size divided by feature_group_count must equal the rhs input feature dimension size, but (.+?) // (.+?) != (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/convolution.py","lineNumber":414,"sourceCode":"  if len(lhs.shape) != len(rhs.shape):\n    msg = (\"conv_general_dilated lhs and rhs must have the same number of \"\n           \"dimensions, but got {} and {}.\")\n    raise ValueError(msg.format(lhs.shape, rhs.shape))\n  if not feature_group_count > 0:\n    msg = (\"conv_general_dilated feature_group_count \"\n           \"must be a positive integer, got {}.\")\n    raise ValueError(msg.format(feature_group_count))\n  lhs_feature_count = lhs.shape[dimension_numbers.lhs_spec[1]]\n  quot, rem = divmod(lhs_feature_count, feature_group_count)\n  if rem:\n    msg = (\"conv_general_dilated feature_group_count must divide lhs feature \"\n           \"dimension size, but {} does not divide {}.\")\n    raise ValueError(msg.format(feature_group_count, lhs_feature_count))\n  if not core.definitely_equal(quot, rhs.shape[dimension_numbers.rhs_spec[1]]):\n    msg = (\"conv_general_dilated lhs feature dimension size divided by \"\n           \"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:","sourceCodeStart":396,"sourceCodeEnd":432,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/convolution.py#L396-L432","documentation":"After dividing the lhs feature dimension by feature_group_count, the result must equal the rhs (kernel) input-feature dimension — each group convolves a slice of size lhs_feature_count/groups against the kernel. If definitely_equal fails, ValueError reports lhs_feature // groups != rhs_in_features.","triggerScenarios":"Grouped conv where kernel input-feature dim doesn't match channels-per-group, e.g. lhs has 8 channels, groups=2 (4 per group) but kernel shaped (H,W,8,Cout) instead of (H,W,4,Cout).","commonSituations":"Reusing an ungrouped kernel shape after adding feature_group_count; PyTorch->JAX ports where weight shape interacts differently with groups (torch divides both in and out channels by groups).","solutions":["Shape the kernel as (..., lhs_in_channels // feature_group_count, out_channels)","Cross-check with PyTorch semantics: torch weight is (out//g, in//g, H, W); JAX rhs is (H, W, in//g, out)","Set feature_group_count=1 if you didn't intend grouping"],"exampleFix":"// before\nk = jax.random.normal(key, (3, 3, 8, 16))  # 8 == full in-channels\nlax.conv_general_dilated(x_c8, k, (1,1), 'SAME', feature_group_count=2)\n// after\nk = jax.random.normal(key, (3, 3, 4, 16))  # 4 == 8 // 2\nlax.conv_general_dilated(x_c8, k, (1,1), 'SAME', feature_group_count=2)","handlingStrategy":"validation","validationCode":"per_group = lhs_in_channels // feature_group_count\nassert kernel.shape[rhs_in_feature_axis] == per_group, (kernel.shape, per_group)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Build kernels as (..., in_ch // groups, out_ch)","Translate torch grouped-conv weight shapes carefully"],"tags":["jax","lax","convolution","feature-group-count","kernel-shape"],"backgroundTag":"group-count-not-divisible","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}