jax-ml/jax · error · ValueError

conv_general_dilated window and window_strides must have the

Error message

conv_general_dilated window and window_strides must have the same number of dimensions, but got {} and {}

What it means

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.

Source

Thrown at jax/_src/lax/convolution.py:446

    msg = ("conv_general_dilated batch_group_count must divide lhs batch "
           "dimension size, but {} does not divide {}.")
    raise ValueError(msg.format(batch_group_count, lhs_batch_count))

  if rhs.shape[dimension_numbers.rhs_spec[0]] % batch_group_count:
    msg = ("conv_general_dilated rhs output feature dimension size must be a "
           "multiple of batch_group_count, but {} is not a multiple of {}.")
    raise ValueError(msg.format(rhs.shape[dimension_numbers.rhs_spec[0]],
                                batch_group_count))

  if batch_group_count > 1 and feature_group_count > 1:
    msg = ("At most one of batch_group_count and feature_group_count may be > "
           "1, got batch_group_count={} and feature_group_count={}")
    raise ValueError(msg.format(batch_group_count, feature_group_count))

  if len(_conv_sdims(dimension_numbers.rhs_spec)) != len(window_strides):
    msg = ("conv_general_dilated window and window_strides must have "
           "the same number of dimensions, but got {} and {}")
    raise ValueError(
        msg.format(len(_conv_sdims(dimension_numbers.rhs_spec)), len(window_strides)))

  lhs_perm, rhs_perm, out_perm = dimension_numbers
  lhs_trans = lax._dilate_shape(np.take(lhs.shape, lhs_perm), lhs_dilation)
  rhs_trans = lax._dilate_shape(np.take(rhs.shape, rhs_perm), rhs_dilation)
  out_trans = conv_shape_tuple(lhs_trans, rhs_trans, window_strides, padding,
                               batch_group_count)
  return tuple(np.take(out_trans, np.argsort(out_perm)))


def _conv_general_dilated_sharding_rule(
    lhs: core.ShapedArray, rhs: core.ShapedArray, *, window_strides, padding,
    lhs_dilation, rhs_dilation, dimension_numbers, feature_group_count,
    batch_group_count, out_sharding, **unused_kwargs):
  if out_sharding is not None:
    assert isinstance(out_sharding, NamedSharding)
    return out_sharding
  # Only allow if rhs is fully replicated and lhs's feature dim is not sharded

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Set window_strides length = number of spatial dims (len(shape)-2, or len(_conv_sdims(rhs_spec)))
2. Validate dimension_numbers matches both tensor ranks
3. Use lax.conv_dimension_numbers to build the spec correctly

Example fix

# before
out = 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')))
# after
out = 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')))

n_spatial = len(lhs.shape) - 2
assert len(window_strides) == n_spatial, (window_strides, lhs.shape)

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• Wrong number of strides for spatial dimensions
• conv_general_dilated batch_group_count must divide lhs batch
• conv_general_dilated rhs output feature dimension size must
• Wrong number of explicit pads for convolution: expected {},
• Negative padding is larger than the size of the correspondin

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/49a5d9c316ea232d. Report an issue: GitHub.