jax-ml/jax · error · TypeError

Symbolic dimension {operation_name} not supported for {p}.

Error message

Symbolic dimension {operation_name} not supported for {p}.

What it means

Internal coercion turning a value into a symbolic dimension polynomial accepts only existing _DimExpr (same scope) or integers. Anything else — float, None, string, numpy float scalar — cannot become a dim and raises TypeError naming the operation.

Source

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

core.pytype_aval_mappings[_DimExpr] = _DimExpr._get_aval
dtypes.register_weak_scalar_type(_DimExpr)

def _convertible_to_int(p: Any) -> TypeGuard[SupportsIndex]:
  try:
    op.index(p)
    return True
  except:
    return False

def _ensure_poly(p: DimSize,
                 operation_name: str,
                 scope: SymbolicScope) -> _DimExpr:
  if isinstance(p, _DimExpr):
    scope._check_same_scope(p, when=f"for operation {operation_name}")
    return p
  if _convertible_to_int(p):
    return _DimExpr(((_DimTerm_one, op.index(p)),), scope)
  raise TypeError(f"Symbolic dimension {operation_name} not supported for {p}.")

def _convertible_to_poly(p: Any) -> bool:
  return isinstance(p, _DimExpr) or _convertible_to_int(p)

def is_symbolic_dim(p: DimSize) -> TypeGuard[_DimExpr]:
  """Checks if a dimension is symbolic.
  """
  return isinstance(p, _DimExpr)


def symbolic_dim_bounds(
    dimension: DimSize | SupportsIndex,
) -> tuple[float, float]:
  """Returns inclusive bounds that JAX can prove for a dimension expression.

  The returned bounds are conservative and may not be tight. Infinite bounds
  mean that JAX could not establish a finite bound; they do not prove that the
  dimension is unbounded.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use integer constants in shape arithmetic (int(2.0) or //2 instead of *0.5)
2. Check for None before shape computations
3. Keep all shape arithmetic in ints or existing symbolic dims

Example fix

# before
new_dim = old_dim * 0.5
# after
new_dim = old_dim // 2  # integer shape arithmetic

assert isinstance(v, (int,)) or is_symbolic_dim(v), f'bad operand in shape math: {v!r}'

def dim_arith_ok(p) -> bool:
    return isinstance(p, int) or is_symbolic_dim(p)

try:
    d = dim * factor
except TypeError as e:
    if 'not supported' in str(e): d = dim * int(factor)

When it happens

Trigger: Arithmetic like dim * 2.0, dim + np.float32(1), or None propagating into shape arithmetic during polymorphic export; operations named in the message (e.g. 'mul', 'parse_constraint').

Common situations: Mixing float scalars (e.g. mean factors like 0.5 * dim for downsampling) with dimension expressions; None from an optional config leaking into shape math.

Related errors

• Encountered unexpected shape dimension {d}
• {name} was requested to map a value of non-array type {core.
• primal and tangent arguments to jax.jvp must be tuples or li
• {prim_name} takes a scalar pred as argument, got {pred}
• `compute_on`'s compute_type argument must be a string.

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