jax-ml/jax · error · ValueError

Error refining shapes. {dump_module_message(module, "before_

Error message

Error refining shapes. {dump_module_message(module, "before_refine_polymorphic_shapes")}

What it means

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__`).

Source

Thrown at jax/_src/interpreters/mlir.py:3560

          for r, aval in zip(rw.results, ctx.avals_out)]


def refine_polymorphic_shapes(module: ir.Module) -> ir.Module:
  """Refines the polymorphic shapes inside a module.

  Given a module with static input shapes, but using dynamic shapes due to
  shape polymorphism, runs shape refinement to resolve all the dynamic shapes.
  Then verifies that there are no more dynamic shapes in the module.
  """
  try:
    refine_polymorphic_shapes = partial(_jax.mlir.refine_polymorphic_shapes,
            mlir_module=module_to_bytecode(module),
            enable_shape_assertions=True,
            validate_static_shapes=True)
    refined_module_str = refine_polymorphic_shapes(
        enable_shardy=config.use_shardy_partitioner.value)
  except Exception as e:
    raise ValueError(
        "Error refining shapes. " +
        dump_module_message(module, "before_refine_polymorphic_shapes")) from e

  context = make_ir_context()
  with context:
    return ir.Module.parse(refined_module_str)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Read the chained exception (raise ... from e) — the true cause is in `e.__cause__` and the module dump shows the failing IR
2. Simplify the symbolic shape expressions (avoid unusual dimension arithmetic) or pin exact shapes where possible
3. If the module dump shows a specific op failing, restructure the computation around that op
4. Search/file the JAX issue tracker with the module dump if it looks like an internal pass bug

Example fix

# before
def f(x):  # x: f32[b, b+1]
    return x[:, 1:]

# after
def f(x):  # keep dimension expressions simple and provably consistent
    return x[:, : x.shape[1] - 1]

try:
    exported = jax.export.export(jit(f))(*args)
except ValueError as e:
    if 'Error refining shapes' in str(e):
        cause = e.__cause__  # real refinement failure
        log_module_dump(str(e))  # contains the IR dump

When it happens

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

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

Related errors

• Cannot lower effectful `reduce_window`.
• Too many dynamic shapes in the input. Mosaic currently only
• multi-platform lowering for buffer_callback
• Nesting `compute_on` with different compute types is not all
• Sharding rule has {len(rule.operand_mappings)} operands, but

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