jax-ml/jax · error · ValueError

Invalid mixing of symbolic scopes {when}.\nExpected {self_de

Error message

Invalid mixing of symbolic scopes {when}.\nExpected {self_descr}scope {self}\nand found for '{other}' ({other_descr}) scope {other.scope}\nSee https://docs.jax.dev/en/latest/export/shape_poly.html#user-specified-symbolic-constraints.

What it means

All symbolic dimension expressions in one computation must belong to the same SymbolicScope (the constraint context created per export). Mixing dims from different scopes — e.g. reusing a variable from an earlier export or a manually built scope — makes constraint reasoning ill-defined and raises ValueError.

Source

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

            _DimTerm.from_operation(
                _DimFactor.FLOORDIV, *before_f.operands, scope=constr.e1.scope),
            before_k2, scope=constr.e1.scope)
        before_e1 = before_f.operands[0]
        self._process_explicit_constraint(
            _SymbolicConstraint(cmp=Comparator.EQ,
                                e1=k_times_floordiv, e2=before_e1,
                                diff=k_times_floordiv - before_e1,
                                debug_str=f"{k_times_floordiv} == {before_e1}")
        )

    self._explicit_constraints.append(constr)

  def _check_same_scope(self, other: _DimExpr,
                        when: str = "",
                        self_descr: str = " ",
                        other_descr: str = "unknown"):
    if self is not other.scope:
      raise ValueError(
          f"Invalid mixing of symbolic scopes {when}.\n"
          f"Expected {self_descr}scope {self}\n"
          f"and found for '{other}' ({other_descr}) scope {other.scope}\n"
          f"See https://docs.jax.dev/en/latest/export/shape_poly.html#user-specified-symbolic-constraints.")

  def _clear_caches(self):
    self._bounds_cache.clear()


class BoundsPrecision(enum.Enum):
  """Specifies desired precision for the bounds calculation.

  Since the bounds calculations are expensive, we allow the caller to specify
  a sufficient condition for a result. As the bounds calculations progresses
  the lower bounds is progressively increased and the upper bounds is
  progressively decreased. Depending on the precision, we may stop the
  computation early, if the results are sufficient for the use case.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Re-export the functions together so they share one scope, or pass shapes as strings so each export creates its own consistent scope
Avoid reusing dim expressions or scope objects across different exports
Ensure the same polymorphic_shapes strings (and constraints) are used consistently when composing exported functions

Example fix

# before
exp1 = jax.export.export(shapes('n,'))(f)
exp2 = jax.export.export(shapes('m,'))(g)  # different scope vars mixed later
# after
# use the same variable name and a single export or shared scope
exp2 = jax.export.export(shapes('n,'), constraints=('n >= 1',))(g)

Defensive patterns

Strategy: validation

Validate before calling

# verify all dim exprs share one scope before combining
assert len({d.scope for d in dims}) == 1, 'mixed symbolic scopes'

Type guard

def same_scope(*dims) -> bool:
    scopes = {getattr(d, 'scope', None) for d in dims}
    return len(scopes) == 1

Try / catch

try:
    out = compose(exp1, exp2)
except ValueError as e:
    if 'mixing of symbolic scopes' in str(e): reexport_together()

Prevention

Never reuse dim expressions or scopes across separate exports
Use identical shape variable names and constraints when composing exports

When it happens

Trigger: Calling a previously exported polymorphic function with arguments whose shapes contain dim vars from another SymbolicScope; combining dims from two different exported functions; capturing a global scope variable across exports.

Common situations: Caching an exported function and composing it with another export; reusing shape objects or scope globals between two jax.export calls; pickling an export and mixing with fresh scopes.

Related errors

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