jax-ml/jax · error · UnexpectedDimVar

Encountered dimension variable

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.

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.

Appendix: 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)