{"record":{"id":"c484e2d3b40ee6b6","repo":"jax-ml/jax","slug":"error-refining-shapes-dump-module-message-module","errorCode":null,"errorMessage":"Error refining shapes. {dump_module_message(module, \"before_refine_polymorphic_shapes\")}","messagePattern":"Error refining shapes\\. (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/interpreters/mlir.py","lineNumber":3560,"sourceCode":"          for r, aval in zip(rw.results, ctx.avals_out)]\n\n\ndef refine_polymorphic_shapes(module: ir.Module) -> ir.Module:\n  \"\"\"Refines the polymorphic shapes inside a module.\n\n  Given a module with static input shapes, but using dynamic shapes due to\n  shape polymorphism, runs shape refinement to resolve all the dynamic shapes.\n  Then verifies that there are no more dynamic shapes in the module.\n  \"\"\"\n  try:\n    refine_polymorphic_shapes = partial(_jax.mlir.refine_polymorphic_shapes,\n            mlir_module=module_to_bytecode(module),\n            enable_shape_assertions=True,\n            validate_static_shapes=True)\n    refined_module_str = refine_polymorphic_shapes(\n        enable_shardy=config.use_shardy_partitioner.value)\n  except Exception as e:\n    raise ValueError(\n        \"Error refining shapes. \" +\n        dump_module_message(module, \"before_refine_polymorphic_shapes\")) from e\n\n  context = make_ir_context()\n  with context:\n    return ir.Module.parse(refined_module_str)\n","sourceCodeStart":3542,"sourceCodeEnd":3567,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/interpreters/mlir.py#L3542-L3567","documentation":"During lowering, JAX runs a pass that refines polymorphic shapes (from jax2static / shape polymorphism) and inserts shape assertions. If that pass throws for any reason, JAX wraps the failure and dumps the module so you can inspect the IR that failed. The underlying exception is chained (see `__cause__`).","triggerScenarios":"Using dynamic shapes (polymorphic dimensions, e.g. jax.export or shapes with symbolic bounds) where the refinement pass cannot prove consistency — e.g. mismatched symbolic dimension expressions, invalid rewrites, or a bug in shape refinement for a particular op.","commonSituations":"jax.export with symbolic dimensions; automatic-shape-refinement regressions after a JAX upgrade; expressions the Shardy/refinement pass cannot handle. Inspect the printed 'before_refine_polymorphic_shapes' module dump and the chained cause.","solutions":["Read the chained exception (raise ... from e) — the true cause is in `e.__cause__` and the module dump shows the failing IR","Simplify the symbolic shape expressions (avoid unusual dimension arithmetic) or pin exact shapes where possible","If the module dump shows a specific op failing, restructure the computation around that op","Search/file the JAX issue tracker with the module dump if it looks like an internal pass bug"],"exampleFix":"# before\ndef f(x):  # x: f32[b, b+1]\n    return x[:, 1:]\n\n# after\ndef f(x):  # keep dimension expressions simple and provably consistent\n    return x[:, : x.shape[1] - 1]","handlingStrategy":"try-catch","validationCode":null,"typeGuard":null,"tryCatchPattern":"try:\n    exported = jax.export.export(jit(f))(*args)\nexcept ValueError as e:\n    if 'Error refining shapes' in str(e):\n        cause = e.__cause__  # real refinement failure\n        log_module_dump(str(e))  # contains the IR dump","preventionTips":["Keep symbolic dimension expressions simple and consistent","Pin exact shapes during development, then relax to polymorphic ones","Save the dumped module when reporting shape-refinement issues"],"tags":["jax","polymorphic-shapes","dynamic-shapes","mlir","lowering"],"backgroundTag":"shape-refinement-failed","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}