{"record":{"id":"49a5d9c316ea232d","repo":"jax-ml/jax","slug":"conv-general-dilated-window-and-window-strides-mus","errorCode":null,"errorMessage":"conv_general_dilated window and window_strides must have the same number of dimensions, but got {} and {}","messagePattern":"conv_general_dilated window and window_strides 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":446,"sourceCode":"    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,\n    batch_group_count, out_sharding, **unused_kwargs):\n  if out_sharding is not None:\n    assert isinstance(out_sharding, NamedSharding)\n    return out_sharding\n  # Only allow if rhs is fully replicated and lhs's feature dim is not sharded","sourceCodeStart":428,"sourceCodeEnd":464,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/convolution.py#L428-L464","documentation":"window_strides must have one entry per spatial dimension of the convolution. The rule compares len(window_strides) with the number of spatial dims implied by the rhs spec in dimension_numbers. Mismatched lengths raise this error.","triggerScenarios":"Calling jax.lax.conv_general_dilated with a 3D conv (3 spatial dims) but window_strides=(1,1), or vice versa; also mismatched ConvDimensionNumbers where the rhs spec implies a different spatial rank.","commonSituations":"Reusing stride tuples from 2D code in 3D convs; constructing dimension_numbers by hand with the wrong length; porting between NHWC/NCHW layouts with different rank tensors.","solutions":["Set window_strides length = number of spatial dims (len(shape)-2, or len(_conv_sdims(rhs_spec)))","Validate dimension_numbers matches both tensor ranks","Use lax.conv_dimension_numbers to build the spec correctly"],"exampleFix":"# before\nout = lax.conv_general_dilated(lhs_5d, rhs_5d, (1, 1), 'VALID', dimension_numbers=lax.conv_dimension_numbers(lhs_5d.shape, rhs_5d.shape, ('NCDHW', 'OCIHW', 'NCDHW')))\n# after\nout = lax.conv_general_dilated(lhs_5d, rhs_5d, (1, 1, 1), 'VALID', dimension_numbers=lax.conv_dimension_numbers(lhs_5d.shape, rhs_5d.shape, ('NCDHW', 'OCIHW', 'NCDHW')))","handlingStrategy":"validation","validationCode":"n_spatial = len(lhs.shape) - 2\nassert len(window_strides) == n_spatial, (window_strides, lhs.shape)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Derive strides/dilations/padding tuples from len(shape)-2 instead of hardcoding","Build dimension_numbers with lax.conv_dimension_numbers to keep rank consistent"],"tags":["jax","convolution","strides","shape-validation"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}