[shape_poly] Fix the handling of __pow__ for symbolic dimensions

The code for handling exponentiation was wrong, and there were
no tests.

Author: gnecula

jax/_src/export/shape_poly.py

@@ -764,13 +764,23 @@ def __rmul__(self, other):
       return _DimExpr._linear_combination(self, other, 0, 0, self.scope)
     return _ensure_poly(other, "mul", self.scope).__mul__(self)
 
-  def __pow__(self, power, modulo=None):
-    assert modulo is None
-    try:
-      power = int(power)
-    except:
-      raise InconclusiveDimensionOperation(f"Symbolic dimension cannot be raised to non-integer power '{self}' ^ '{power}'")
-    return functools.reduce(op.mul, [self] * power)
+  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)  # type: ignore
+    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):

tests/shape_poly_test.py

@@ -128,11 +128,11 @@ def sampled_assertion(self,
                         ):
     """Checks `assertion(e, fun(*operands))` symbolically and concretely.
 
-    For the concrete check, it will same the space of dimension variable
+    For the concrete check, it will sample the space of dimension variable
     assignments for the dimension variables in `e`.
 
-    This is useful when `fun` can operate both with polynomials and with
-    concrete values, and we want to double-check that the behavior is sound.
+    This is useful when `fun` can operate both with symbolic and with
+    concrete values, and we want to check that the behavior is sound.
     """
     computed_sym = fun(*operands_sym)
     assertion_fun = {
@@ -1429,6 +1429,29 @@ def test_non_trivial_dim_expr(self, expr=lambda d: d % -2):
       arg_descriptors=[RandArg((3,), np.int64)],
       polymorphic_shapes=["b"])
 
+  @jtu.parameterized_filterable(
+      # The function `f` will be called with x: f32[b]
+      kwargs=[
+          dict(testcase_name="cube", f=lambda x: x.shape[0] ** 3),
+          dict(testcase_name="zero", f=lambda x: x.shape[0] ** 0),
+          dict(testcase_name="rpow", f=lambda x: 2 ** x.shape[0]),
+          dict(testcase_name="negative",
+               f=lambda x: x.shape[0] ** -2,
+               expect_error=(ValueError, "cannot be raised to negative powers")),
+          dict(testcase_name="non_integer",
+               f=lambda x: x.shape[0] ** 1.5,
+               expect_error=(ValueError, "cannot be raised to non-integer powers")),
+          dict(testcase_name="sym_pow",
+               f=lambda x: x.shape[0] ** x.shape[0]),
+      ]
+  )
+  def test_pow(self, f, expect_error: tuple[Exception, str] | None = None):
+    check_shape_poly(self,
+                     f,
+                     arg_descriptors=[RandArg((3,), np.float32)],
+                     polymorphic_shapes=["b"],
+                     expect_error=expect_error)
+
   def test_static_shape_result(self):
     """The result has static shape."""