{"record":{"id":"2791e84d155d50d7","repo":"jax-ml/jax","slug":"convolutions-with-non-static-strides-dilation-fe","errorCode":null,"errorMessage":"Convolutions with non-static strides, dilation, feature_group_count, or batch_group_count","messagePattern":"Convolutions with non-static strides, dilation, feature_group_count, or batch_group_count","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/convolution.py","lineNumber":823,"sourceCode":"    input_feature_dimension=lhs_spec[1],\n    input_spatial_dimensions=list(lhs_spec[2:]),\n    kernel_output_feature_dimension=rhs_spec[0],\n    kernel_input_feature_dimension=rhs_spec[1],\n    kernel_spatial_dimensions=list(rhs_spec[2:]),\n    output_batch_dimension=out_spec[0],\n    output_feature_dimension=out_spec[1],\n    output_spatial_dimensions=list(out_spec[2:]))\n  num_spatial_dims = len(rhs_spec) - 2\n  if len(padding) == 0:\n    padding = np.zeros((0, 2), dtype=np.int64)\n  window_reversal = ir.DenseBoolArrayAttr.get([False] * num_spatial_dims)\n  if (not core.is_constant_shape(window_strides) or\n      not core.is_constant_shape(lhs_dilation) or\n      not core.is_constant_shape(rhs_dilation) or\n      not core.is_constant_dim(feature_group_count) or\n      not core.is_constant_dim(batch_group_count)):\n    # TODO(https://github.com/openxla/stablehlo/issues/1268)\n    raise NotImplementedError(\"Convolutions with non-static strides, dilation, feature_group_count, or batch_group_count\")\n  if all(core.is_constant_shape(p) for p in padding):\n    result_type = mlir.aval_to_ir_type(ctx.module_context, aval_out)\n    out = hlo.convolution(\n        result_type, lhs, rhs,\n        dimension_numbers=dnums,\n        feature_group_count=mlir.i64_attr(feature_group_count),\n        batch_group_count=mlir.i64_attr(batch_group_count),\n        window_strides=mlir.dense_int_array(window_strides),\n        padding=mlir.dense_int_elements(padding),\n        lhs_dilation=mlir.dense_int_array(lhs_dilation),\n        rhs_dilation=mlir.dense_int_array(rhs_dilation),\n        window_reversal=window_reversal,\n        precision_config=lax.precision_attr(precision))\n    return [mlir.lower_with_sharding_in_types(ctx, out, aval_out)]\n  else:\n    # d_padding will be an array i32[N, 2] with pad_lo and pad_hi for each\n    # spatial dimension.\n    int2d = mlir.aval_to_ir_type(ctx.module_context, core.ShapedArray((1, 2), np.int32))","sourceCodeStart":805,"sourceCodeEnd":841,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/convolution.py#L805-L841","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Mark the values as static: pass them as Python ints or close over constants, or use static_argnums/static_argnames in jit","Restructure so strides/groups are computed outside the traced function","For truly dynamic conv config, dispatch to different compiled specializations keyed on the static value"],"exampleFix":"# before\n@jax.jit\ndef f(x, w, strides):\n    return lax.conv_general_dilated(x, w, strides, 'SAME')\n# after\n@jax.jit\ndef f(x, w):\n    return lax.conv_general_dilated(x, w, (1, 1), 'SAME')","handlingStrategy":"validation","validationCode":"assert not isinstance(window_strides, jax.Array) or window_strides.is_constant() if hasattr(window_strides, 'is_constant') else True\n# simplest: ensure plain python ints\ndef check_static(seq):\n    return all(isinstance(v, (int, np.integer)) for v in seq)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Pass strides/dilations/group counts as Python ints, never traced values","Use jax.jit(static_argnames=(...)) for config arguments"],"tags":["jax","jit","static-shapes","convolution","lowering"],"backgroundTag":"non-static-traced-value","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}