jax-ml/jax · error · UnexpectedDimVar

Encountered dimension variable '{self.var}' that is not appe

Error message

Encountered dimension variable '{self.var}' that is not appearing in the shapes of the function arguments.\nThe following dimension variables are appearing in the shapes of the function arguments: {list(env.keys())}.\nPlease see https://docs.jax.dev/en/latest/export/shape_poly.html#dimension-variables-must-be-solvable-from-the-input-shapes for more details.

What it means

During polymorphic export evaluation, JAX solved dimension expressions against the shapes of the concrete (or specified) arguments. A dimension variable like 'n' appeared in an intermediate shape but never in any input argument shape, so its value cannot be solved.

Source

Thrown at jax/_src/export/shape_poly.py:230

  def evaluate(self, env: DimVarEnv, scope: SymbolicScope):
    from jax._src.lax import lax

    if self.var is not None:
      try:
        return env[self.var]
      except KeyError:
        # Perhaps there is a normalization rule for this variable
        normalized_var = _DimExpr._from_var(self.var, scope)
        if core.is_constant_dim(normalized_var):
          return normalized_var
        non_trivial_normalization = (v1 := normalized_var._to_var()) is None or v1 != self.var  # pyrefly: ignore[missing-attribute]
        if non_trivial_normalization:
          return normalized_var._evaluate(env)  # pyrefly: ignore[missing-attribute]
        err_msg = (
            f"Encountered dimension variable '{self.var}' that is not appearing in the shapes of the function arguments.\n"
            f"The following dimension variables are appearing in the shapes of the function arguments: {list(env.keys())}.\n"
            "Please see https://docs.jax.dev/en/latest/export/shape_poly.html#dimension-variables-must-be-solvable-from-the-input-shapes for more details.")
        raise UnexpectedDimVar(err_msg)
    else:
      operand_values = [opnd._evaluate(env) for opnd in self.operands]
      if self.operation == _DimFactor.FLOORDIV:
        return divmod(*operand_values)[0]
      elif self.operation == _DimFactor.MOD:
        return divmod(*operand_values)[1]
      elif self.operation == _DimFactor.MAX:
        op1, op2 = operand_values
        if core.is_constant_dim(op1) and core.is_constant_dim(op2):
          return max(op1, op2)
        if core.is_symbolic_dim(op1) or core.is_symbolic_dim(op2):
          return core.max_dim(op1, op2)
        # In the context of `evaluate` dimension variables may be mapped to
        # JAX Tracers.
        return lax.max(op1, op2)
      elif self.operation == _DimFactor.MIN:
        op1, op2 = operand_values
        if core.is_constant_dim(op1) and core.is_constant_dim(op2):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Make the variable appear in at least one input argument's shape (e.g. add a dummy argument of shape (n,) or restructure so the var is derived from an input)
Use a constant instead of the unsolvable variable for that dimension
Rework polymorphic_shapes so all internal dim vars are expressible from input shapes; see linked shape_poly docs section

Example fix

# before
def f(x):  # x: (m,)
    return jnp.arange(x.shape[0] * 2)
jax.export.export(jax.export.shapes('n,'))(f)(x)  # n never in args
# after
def f(x, hint):  # hint: (n,)
    return jnp.arange(x.shape[0] * 2)
jax.export.export(jax.export.shapes('n,n,'))(f)(x, jnp.zeros(n))

Defensive patterns

Strategy: validation

Validate before calling

# ensure every dim var used internally appears in an input shape
import jax.export as jex
used = collect_dim_vars(fn)  # e.g. by test-tracing with symbolic shapes
provided = set(''.join(c for c in s if c.isalpha()) for s in polymorphic_shapes)
assert used <= provided, f'unsolvable vars: {used - provided}'

Try / catch

try:
    exp = jax.export.export(shapes)(fn)
except jax.export.shape_poly.UnexpectedDimVar as e:
    add_dummy_input_exposing_var(str(e))

Prevention

Parameterize every internal quantity as a function of input-shape variables
Write a smoke export test with symbolic shapes before production runs

When it happens

Trigger: Exporting with polymorphic shapes where an equation produces a dim var (e.g. via random epochs or reshape) that is not present in any input dimension, e.g. shape '(n, n*2)' but computing something like a dimension derived from an internal constant divisible by a var.

Common situations: Using jax.export with shape polymorphism where a variable only appears on an output or inside jnp.arange, or specifying polymorphic_shapes strings that omit a variable used internally (e.g. random-wrapper epochs defaulting to a symbolic var).

Related errors

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