{"record":{"id":"2b9a6ea5f0599e33","repo":"jax-ml/jax","slug":"conv-general-dilated-lhs-and-rhs-must-have-the-sam","errorCode":null,"errorMessage":"conv_general_dilated lhs and rhs must have the same number of dimensions, but got {} and {}.","messagePattern":"conv_general_dilated lhs and rhs must have the same number of dimensions, but got (.+?) and (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/convolution.py","lineNumber":399,"sourceCode":"      pads = padding\n  if transpose_kernel:\n    # flip spatial dims and swap input / output channel axes\n    rhs = _flip_axes(rhs, np.array(dn.rhs_spec)[2:])\n    rhs = rhs.swapaxes(dn.rhs_spec[0], dn.rhs_spec[1])\n  return conv_general_dilated(lhs, rhs, one, pads, strides, rhs_dilation, dn,\n                              precision=precision,\n                              preferred_element_type=preferred_element_type)\n\n\ndef _conv_general_dilated_shape_rule(\n    lhs: core.ShapedArray, rhs: core.ShapedArray, *, window_strides, padding,\n    lhs_dilation, rhs_dilation, dimension_numbers, feature_group_count,\n    batch_group_count, **unused_kwargs) -> tuple[int, ...]:\n  assert type(dimension_numbers) is ConvDimensionNumbers\n  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 \"","sourceCodeStart":381,"sourceCodeEnd":417,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/convolution.py#L381-L417","documentation":"The shape rule for conv_general_dilated requires lhs (input) and rhs (kernel) to have identical rank, since a single ConvDimensionNumbers spec maps dimensions of both. If len(lhs.shape) != len(rhs.shape) it raises ValueError showing both shapes.","triggerScenarios":"Convolving a 4-D input (NHWC) with a 3-D kernel (HWI) or similar rank mismatch in lax.conv/conv_general_dilated; often via the Python-level wrapper passing unbatched kernels.","commonSituations":"Forgetting to add the batch dimension (passing HWC input against HWIO kernel); sharing weights across a batch but reshaping the kernel incorrectly; vmap or tree_map stripping a leading dim from only one operand.","solutions":["Add/remove leading dimensions to match ranks: x[None] for the batch dim, or kernel as (H,W,Cin,Cout)","Check both .ndim before calling: assert lhs.ndim == rhs.ndim","Use convenience wrappers (jnp.convolution) that handle common ranks"],"exampleFix":"// before\nout = lax.conv(x_hwc, kernel_hwio, (1,1), 'SAME')  # ranks 3 vs 4\n// after\nout = lax.conv(x_nhwc, kernel_hwio, (1,1), 'SAME')  # x[None] added","handlingStrategy":"validation","validationCode":"assert lhs.ndim == rhs.ndim, (lhs.shape, rhs.shape)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Add x[None] batch dim before lax.conv","vmap instead of manual rank gymnastics"],"tags":["jax","lax","convolution","shape-mismatch"],"backgroundTag":"operand-rank-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}