{"record":{"id":"dc539c55fc9865b0","repo":"jax-ml/jax","slug":"pow-modulo-not-implemented","errorCode":null,"errorMessage":"__pow__ modulo not implemented","messagePattern":"__pow__ modulo not implemented","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/export/shape_poly.py","lineNumber":763,"sourceCode":"    for mon1, coeff1 in self._sorted_terms:\n      for mon2, coeff2 in other._sorted_terms:\n        mon = mon1.mul(mon2)\n        _DimExpr._add_coeff(coeffs, mon, coeff1 * coeff2)\n    return _DimExpr._normalize_sorted_terms(_DimExpr._coeff_to_sorted_terms(coeffs),\n                                            self.scope)\n\n  def __rmul__(self, other):\n    if isinstance(other, core.Tracer) or not _convertible_to_poly(other):\n      return self.__jax_array__().__rmul__(other)\n    if isinstance(other, int):\n      if other == 1: return self\n      if other == 0: return 0\n      return _DimExpr._linear_combination(self, other, 0, 0, self.scope)\n    return _ensure_poly(other, \"mul\", self.scope).__mul__(self)\n\n  def __pow__(self, power: core.DimSize, modulo=None):\n    if modulo is not None:\n      raise NotImplementedError(\"__pow__ modulo not implemented\")\n    if is_symbolic_dim(power):\n      return power.__rpow__(self)\n    if power != int(power):\n      raise ValueError(f\"Symbolic dimension cannot be raised to non-integer powers: '{self}' ** '{power}'\")\n    if power >= 0:\n      return functools.reduce(op.mul, [self] * power, 1)\n    # We don't support negative powers, because JAX does not allow negative\n    # powers for integers\n    raise ValueError(f\"Symbolic dimension cannot be raised to negative powers: '{self}' ** '{power}'\")\n\n  def __rpow__(self, other, modulo=None):\n    if modulo is not None:\n      raise NotImplementedError(\"__rpow__ modulo not implemented\")\n    return self.__jax_array__().__rpow__(other)\n\n  def __floordiv__(self, divisor):\n    if isinstance(divisor, core.Tracer) or not _convertible_to_poly(divisor):\n      return self.__jax_array__().__floordiv__(divisor)","sourceCodeStart":745,"sourceCodeEnd":781,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/export/shape_poly.py#L745-L781","documentation":"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.","triggerScenarios":"Calling pow(dim, k, m) or dim ** k % ... via builtins.pow with three args where dim is a symbolic dimension during polymorphic export.","commonSituations":"Rare; typically from generic numeric code paths (e.g. RSA-like arithmetic or hashing utilities) applied to tracer/symbolic dims during export or vmap.","solutions":["Avoid three-argument pow on symbolic dims; compute (dim ** k) then apply % m separately if mathematically valid","Replace the symbolic dim with a concrete integer before the pow-with-modulus","Restructure to use jnp operations on arrays rather than Python scalar pow on dims"],"exampleFix":"# before\nr = pow(dim, 3, 7)\n# after\nr = (dim ** 3) % 7  # if semantics acceptable, else use concrete dim","handlingStrategy":"type-guard","validationCode":"if isinstance(dim, jax.export.shape_poly._DimExpr) and modulo is not None:\n    raise UserError('pow with modulus unsupported on symbolic dims')","typeGuard":"def safe_pow(base, exp, mod=None):\n    if mod is not None and is_symbolic(base): raise TypeError\n    return pow(base, exp, mod)","tryCatchPattern":"try:\n    r = pow(dim, 3, 7)\nexcept NotImplementedError:\n    r = (dim ** 3) % 7","preventionTips":["Never use three-argument pow inside traced/polymorphic code","Keep modular arithmetic on concrete integers or jnp arrays"],"tags":["jax","shape-polymorphism","pow","not-implemented"],"backgroundTag":"unsupported-operator-on-symbolic-dimension","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}