jax-ml/jax · error · NotImplementedError

Convolutions with non-static strides, dilation, feature_grou

Error message

Convolutions with non-static strides, dilation, feature_group_count, or batch_group_count

What it means

When lowering a convolution to StableHLO, JAX requires window_strides, lhs/rhs_dilation, feature_group_count and batch_group_count to be statically known (compile-time constants). Tracer-derived (dynamic) values cannot be represented because StableHLO convolution attributes must be constants.

Source

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

    input_feature_dimension=lhs_spec[1],
    input_spatial_dimensions=list(lhs_spec[2:]),
    kernel_output_feature_dimension=rhs_spec[0],
    kernel_input_feature_dimension=rhs_spec[1],
    kernel_spatial_dimensions=list(rhs_spec[2:]),
    output_batch_dimension=out_spec[0],
    output_feature_dimension=out_spec[1],
    output_spatial_dimensions=list(out_spec[2:]))
  num_spatial_dims = len(rhs_spec) - 2
  if len(padding) == 0:
    padding = np.zeros((0, 2), dtype=np.int64)
  window_reversal = ir.DenseBoolArrayAttr.get([False] * num_spatial_dims)
  if (not core.is_constant_shape(window_strides) or
      not core.is_constant_shape(lhs_dilation) or
      not core.is_constant_shape(rhs_dilation) or
      not core.is_constant_dim(feature_group_count) or
      not core.is_constant_dim(batch_group_count)):
    # TODO(https://github.com/openxla/stablehlo/issues/1268)
    raise NotImplementedError("Convolutions with non-static strides, dilation, feature_group_count, or batch_group_count")
  if all(core.is_constant_shape(p) for p in padding):
    result_type = mlir.aval_to_ir_type(ctx.module_context, aval_out)
    out = hlo.convolution(
        result_type, lhs, rhs,
        dimension_numbers=dnums,
        feature_group_count=mlir.i64_attr(feature_group_count),
        batch_group_count=mlir.i64_attr(batch_group_count),
        window_strides=mlir.dense_int_array(window_strides),
        padding=mlir.dense_int_elements(padding),
        lhs_dilation=mlir.dense_int_array(lhs_dilation),
        rhs_dilation=mlir.dense_int_array(rhs_dilation),
        window_reversal=window_reversal,
        precision_config=lax.precision_attr(precision))
    return [mlir.lower_with_sharding_in_types(ctx, out, aval_out)]
  else:
    # d_padding will be an array i32[N, 2] with pad_lo and pad_hi for each
    # spatial dimension.
    int2d = mlir.aval_to_ir_type(ctx.module_context, core.ShapedArray((1, 2), np.int32))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Mark the values as static: pass them as Python ints or close over constants, or use static_argnums/static_argnames in jit
2. Restructure so strides/groups are computed outside the traced function
3. For truly dynamic conv config, dispatch to different compiled specializations keyed on the static value

Example fix

# before
@jax.jit
def f(x, w, strides):
    return lax.conv_general_dilated(x, w, strides, 'SAME')
# after
@jax.jit
def f(x, w):
    return lax.conv_general_dilated(x, w, (1, 1), 'SAME')

assert not isinstance(window_strides, jax.Array) or window_strides.is_constant() if hasattr(window_strides, 'is_constant') else True
# simplest: ensure plain python ints
def check_static(seq):
    return all(isinstance(v, (int, np.integer)) for v in seq)

When it happens

Trigger: Computing strides, dilations, or group counts from traced values inside jit/vmap/pmap (e.g. strides derived from a shape argument), or using dynamic_slice/shape operations to build them.

Common situations: Writing shape-polymorphic code with jax.experimental.jax2tf or dynamic shapes; deriving kernel strides from input shapes inside a jitted function; passing np arrays computed under a tracer.

Related errors

• multi-platform lowering for buffer_callback
• Formatting arguments to checkify.check need to be PyTrees of
• Nesting `compute_on` with different compute types is not all
• Value of type {type(self)} is not convertible to float.
• Value of type {type(self)} is not convertible to complex.

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