{"record":{"id":"c7203324683c2104","repo":"jax-ml/jax","slug":"shape-polymorphism-is-supported-for-jnp-pad-with","errorCode":null,"errorMessage":"Shape polymorphism is supported for jnp.pad with 'reflect' or 'symmetric' padding mode only when it is possible to determine at lowering time that the axis size (= {axis_size}) is larger than 1 and larger or equal than the padding length (= {padding}). Error while handling {'left' if before else 'right'} padding on axis {i}.","messagePattern":"Shape polymorphism is supported for jnp\\.pad with 'reflect' or 'symmetric' padding mode only when it is possible to determine at lowering time that the axis size \\(= (.+?)\\) is larger than 1 and larger or equal than the padding length \\(= (.+?)\\)\\. Error while handling (.+?) padding on axis (.+?)\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/lax_numpy.py","lineNumber":4017,"sourceCode":"      if before:\n        edge = lax_slicing.slice_in_dim(array, 0, 1, axis=i)\n      else:\n        edge = lax_slicing.slice_in_dim(array, -1, None, axis=i)\n\n      # Try to give nicer error messages for unsupported shape polymorphic uses\n      shape_poly_error_msg = lambda: (\n          \"Shape polymorphism is supported for jnp.pad with 'reflect' or \"\n          \"'symmetric' padding mode only when it is possible to determine \"\n          f\"at lowering time that the axis size (= {axis_size}) is larger than 1 \"\n          f\"and larger or equal than the padding length (= {padding}). \"\n          f\"Error while handling {'left' if before else 'right'} padding on axis {i}.\")\n      try:\n        # We check that we can determine all comparisons.\n        offset = 1 if (mode == \"reflect\" and axis_size > 1) else 0\n        has_poly_dim = not core.is_constant_shape((axis_size, padding))\n        # For shape polymorphism, ensure the loop below ends after 1 iteration\n        if has_poly_dim and not (axis_size > 1 and axis_size - offset >= padding):\n          raise ValueError(shape_poly_error_msg())\n      except core.InconclusiveDimensionOperation as e:\n        raise ValueError(shape_poly_error_msg()) from e\n\n      while padding > 0:\n        curr_pad = min(padding, axis_size - offset)\n        padding -= curr_pad\n        if has_poly_dim: assert padding == 0\n\n        if before:\n          start = offset\n          stop = offset + curr_pad\n        else:\n          start = -(curr_pad + offset)\n          stop = None if (mode == \"symmetric\" or axis_size == 1) else -1\n\n        x = lax_slicing.slice_in_dim(array, start, stop, axis=i)\n        x = flip(x, axis=i)\n","sourceCodeStart":3999,"sourceCodeEnd":4035,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/lax_numpy.py#L3999-L4035","documentation":"With shape polymorphism (symbolic dimensions), jnp.pad's reflect/symmetric implementation must statically decide that the axis size exceeds 1 and is at least the padding length so the padding loop terminates in one iteration. If the symbolic comparison is inconclusive, this ValueError is raised from the explicit check.","triggerScenarios":"Using jax.export/jit with symbolic axis sizes (e.g. batch='n') and calling jnp.pad(..., mode='reflect'|'symmetric') with padding >= the symbolic axis size, or where axis_size > 1 cannot be proven.","commonSituations":"Exporting models with dynamic batch/sequence dimensions that pad with reflection; padding lengths derived from user config larger than the traced dimension.","solutions":["Reduce padding to at most axis_size - 1 (symmetric) / axis_size - 2 (reflect) as a static constant","Use mode='constant' for polymorphic shapes","Constrain the polymorphic dimension with explicit bounds so comparisons are conclusive (e.g. use jax.ShapeDtypeStruct with bounded polys)"],"exampleFix":"// before\njnp.pad(x, ((n - 1, 0),), mode='symmetric')  # n symbolic\n// after\njnp.pad(x, ((2, 0),), mode='symmetric')  # static, small padding","handlingStrategy":"fallback","validationCode":"try:\n    _ = bool(axis_size > 1 and axis_size - 1 >= padding)\nexcept Exception:\n    mode = 'constant'  # cannot verify statically","typeGuard":null,"tryCatchPattern":"try:\n    y = jnp.pad(x, w, mode='reflect')\nexcept ValueError:\n    y = jnp.pad(x, w, mode='constant')","preventionTips":["Keep padding widths literal constants in exported/polymorphic functions","Use constant mode when dimensions are symbolic"],"tags":["jnp-pad","shape-polymorphism","symbolic-dimensions","reflect"],"backgroundTag":"shape-polymorphism-static-shape-required","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}