jax-ml/jax · error · ValueError

Symbolic dimension cannot be raised to non-integer powers: '

Error message

Symbolic dimension cannot be raised to non-integer powers: '{self}' ** '{power}'

What it means

Symbolic dimensions may only be raised to integer powers; the polynomial algebra cannot represent non-integer exponents. A fractional (or otherwise non-integral) power triggers this ValueError.

Source

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

    return _DimExpr._normalize_sorted_terms(_DimExpr._coeff_to_sorted_terms(coeffs),
                                            self.scope)

  def __rmul__(self, other):
    if isinstance(other, core.Tracer) or not _convertible_to_poly(other):
      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):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Ensure the exponent is an int (cast: int(power)) when it is mathematically integral
Restructure to avoid non-integer powers of dims (pass sqrt dimension as a separate symbolic variable)
Drop polymorphism for that submodule (use concrete shapes)

Example fix

# before
side = flat_dim ** 0.5
# after
# parameterize the side length itself in polymorphic_shapes
side = 'h'  # and flat_dim = h * h via constraints or input shapes

Defensive patterns

Strategy: validation

Validate before calling

assert float(power).is_integer(), 'exponent must be integral for symbolic dims'
power = int(power)

Type guard

def integral_pow_ok(p) -> bool:
    return isinstance(p, (int,)) or (isinstance(p, float) and p.is_integer())

Try / catch

try:
    d = dim ** p
except ValueError as e:
    if 'non-integer powers' in str(e): d = dim ** int(p) if float(p).is_integer() else fallback()

Prevention

Cast exponents to int when mathematically integral
Parameterize derived dims (like sqrt sizes) as their own symbolic variables

When it happens

Trigger: dim ** 0.5, dim ** 1.5, or np.float_power applied to a symbolic dim during polymorphic export; also dim ** np.float64(2.0)-like values that compare unequal to their int().

Common situations: Code computing square roots of dimensions (e.g. flattened image dims sqrt(h*w)) under jax.export shape polymorphism; passing numpy float scalars as exponents.

Related errors

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