jax-ml/jax · error · NotImplementedError

__pow__ modulo not implemented

Error message

__pow__ modulo not implemented

What it means

Python's pow protocol allows a third modulo argument (three-argument pow / pow(a, b, m)). JAX's symbolic dimensions only support integer exponentiation without a modulus, so passing modulo raises NotImplementedError.

Source

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

    for mon1, coeff1 in self._sorted_terms:
      for mon2, coeff2 in other._sorted_terms:
        mon = mon1.mul(mon2)
        _DimExpr._add_coeff(coeffs, mon, coeff1 * coeff2)
    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)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Avoid three-argument pow on symbolic dims; compute (dim ** k) then apply % m separately if mathematically valid
2. Replace the symbolic dim with a concrete integer before the pow-with-modulus
3. Restructure to use jnp operations on arrays rather than Python scalar pow on dims

Example fix

# before
r = pow(dim, 3, 7)
# after
r = (dim ** 3) % 7  # if semantics acceptable, else use concrete dim

if isinstance(dim, jax.export.shape_poly._DimExpr) and modulo is not None:
    raise UserError('pow with modulus unsupported on symbolic dims')

def safe_pow(base, exp, mod=None):
    if mod is not None and is_symbolic(base): raise TypeError
    return pow(base, exp, mod)

try:
    r = pow(dim, 3, 7)
except NotImplementedError:
    r = (dim ** 3) % 7

When it happens

Trigger: Calling pow(dim, k, m) or dim ** k % ... via builtins.pow with three args where dim is a symbolic dimension during polymorphic export.

Common situations: Rare; typically from generic numeric code paths (e.g. RSA-like arithmetic or hashing utilities) applied to tracer/symbolic dims during export or vmap.

Related errors

• Symbolic dimension cannot be raised to non-integer powers: '
• Symbolic dimension cannot be raised to negative powers: '{se
• __rpow__ modulo not implemented
• array ref with memory space only works inside of a `jit`.
• pinned array ref only works inside of a `jit`.

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