jax-ml/jax · error · NotImplementedError

__rpow__ modulo not implemented

Error message

__rpow__ modulo not implemented

What it means

Python's reflected-power protocol can also receive a modulo argument; JAX symbolic dimensions do not support modular exponentiation, so __rpow__ with modulo raises NotImplementedError.

Source

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

      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
    return self.__jax_array__().__truediv__(divisor)

  def __rtruediv__(self, dividend):
    # Used for "/", when dividend is not a _DimExpr

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Avoid three-argument pow involving symbolic dims
2. Convert the dim to a concrete int or a jnp array before the pow call
3. Reorder so the symbolic dim is not the exponent of a scalar pow with modulus

Example fix

# before
r = pow(2, dim, 101)
# after
r = jnp.power(2, jnp.asarray(dim)) % 101  # on arrays, outside symbolic dim algebra

if modulo is not None and is_symbolic_dim(exponent): raise TypeError('no modular pow with symbolic dims')

def safe_pow3(b, e, m):
    if m is not None and is_symbolic_dim(e): raise TypeError
    return pow(b, e, m)

try:
    r = pow(2, dim, 101)
except NotImplementedError:
    r = (jnp.power(2, jnp.asarray(dim))) % 101

When it happens

Trigger: pow(const_base, symbolic_dim, modulus) — three-argument pow where the exponent is a symbolic dimension during polymorphic tracing.

Common situations: Very rare; cryptography- or hashing-style code paths executed under jax.export/jit with symbolic dims.

Related errors

• __pow__ modulo not implemented
• array ref with memory space only works inside of a `jit`.
• pinned array ref only works inside of a `jit`.
• Encountered dimension variable '{self.var}' that is not appe
• Cannot divide {self} by {divisor}.

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