jax-ml/jax · error · ValueError

Symbolic dimension cannot be raised to negative powers: '{se

Error message

Symbolic dimension cannot be raised to negative powers: '{self}' ** '{power}'

What it means

JAX's symbolic dimension algebra does not support negative integer powers because JAX itself forbids negative powers on integer types. dim ** -1 therefore raises this ValueError.

Source

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

      return self.__jax_array__().__rmul__(other)
    if isinstance(other, int):
      if other == 1: return self
      if other == 0: return 0
      return _DimExpr._linear_combination(self, other, 0, 0, self.scope)
    return _ensure_poly(other, "mul", self.scope).__mul__(self)

  def __pow__(self, power: core.DimSize, modulo=None):
    if modulo is not None:
      raise NotImplementedError("__pow__ modulo not implemented")
    if is_symbolic_dim(power):
      return power.__rpow__(self)
    if power != int(power):
      raise ValueError(f"Symbolic dimension cannot be raised to non-integer powers: '{self}' ** '{power}'")
    if power >= 0:
      return functools.reduce(op.mul, [self] * power, 1)
    # We don't support negative powers, because JAX does not allow negative
    # powers for integers
    raise ValueError(f"Symbolic dimension cannot be raised to negative powers: '{self}' ** '{power}'")

  def __rpow__(self, other, modulo=None):
    if modulo is not None:
      raise NotImplementedError("__rpow__ modulo not implemented")
    return self.__jax_array__().__rpow__(other)

  def __floordiv__(self, divisor):
    if isinstance(divisor, core.Tracer) or not _convertible_to_poly(divisor):
      return self.__jax_array__().__floordiv__(divisor)
    return self._divmod(divisor)[0]

  def __rfloordiv__(self, other):
    if isinstance(other, core.Tracer) or not _convertible_to_poly(other):
      return self.__jax_array__().__rfloordiv__(other)
    return _ensure_poly(other, "floordiv", self.scope).__floordiv__(self)

  def __truediv__(self, divisor):
    # Used for "/", which always returns a float

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Compute the reciprocal on a float array instead of on the symbolic dim (dim_as_float = jnp.asarray(dim, jnp.float32))
Rewrite 1/dim as division in the calling expression on floats
Avoid exponentiating dims with negative exponents; keep dims for shapes only

Example fix

# before
inv = batch_dim ** -1
# after
inv = 1.0 / jnp.asarray(batch_dim, jnp.float32)

Defensive patterns

Strategy: type-guard

Validate before calling

assert power >= 0 or not is_symbolic_dim(base), 'negative powers unsupported on symbolic dims'

Type guard

def safe_exp(base, p):
    if p < 0 and is_symbolic_dim(base): return 1.0 / jnp.asarray(base, jnp.float32)
    return base ** p

Try / catch

try:
    r = dim ** -1
except ValueError:
    r = 1.0 / jnp.asarray(dim, jnp.float32)

Prevention

Compute reciprocals on float arrays, never on dims
Keep symbolic dims strictly for shape bookkeeping

When it happens

Trigger: dim ** -1 or int ** negative where dim is symbolic (e.g. computing reciprocal counts from a batch dimension) during tracing/export.

Common situations: Generic Python numeric code that computes 1/n as n ** -1 running under polymorphic export or vmap with symbolic dims.

Related errors

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