[array API] implement jnp.pow; alias for jnp.power

docs/jax.numpy.rst

@@ -323,6 +323,7 @@ namespace; they are listed below.
     polysub
     polyval
     positive
+    pow
     power
     printoptions
     prod

jax/_src/numpy/ufuncs.py

@@ -355,6 +355,9 @@ def power(x1: ArrayLike, x2: ArrayLike, /) -> Array:
   # Handle cases #2 and #3 under a jit:
   return _power(x1, x2)
 
+# Array API alias
+pow = power
+
 @partial(jit, inline=True)
 def _power(x1: ArrayLike, x2: ArrayLike) -> Array:
   x1, x2 = promote_shapes("power", x1, x2)  # not dtypes

jax/experimental/array_api/_elementwise_functions.py

@@ -311,7 +311,7 @@ def positive(x, /):
 def pow(x1, x2, /):
   """Calculates an implementation-dependent approximation of exponentiation by raising each element x1_i (the base) of the input array x1 to the power of x2_i (the exponent), where x2_i is the corresponding element of the input array x2."""
   x1, x2 = _promote_dtypes("pow", x1, x2)
-  return jax.numpy.power(x1, x2)
+  return jax.numpy.pow(x1, x2)
 
 
 def real(x, /):

jax/numpy/__init__.py

@@ -403,6 +403,7 @@
     nextafter as nextafter,
     not_equal as not_equal,
     positive as positive,
+    pow as pow,
     power as power,
     rad2deg as rad2deg,
     radians as radians,

jax/numpy/__init__.pyi

@@ -630,6 +630,7 @@ def polymul(a1: ArrayLike, a2: ArrayLike, *, trim_leading_zeros: bool = ...) ->
 def polysub(a1: Array, a2: Array) -> Array: ...
 def polyval(p: Array, x: Array, *, unroll: int = ...) -> Array: ...
 def positive(x: ArrayLike, /) -> Array: ...
+def pow(x: ArrayLike, y: ArrayLike, /) -> Array: ...
 def power(x: ArrayLike, y: ArrayLike, /) -> Array: ...
 printoptions = _np.printoptions
 def prod(a: ArrayLike, axis: _Axis = ..., dtype: DTypeLike = ...,

tests/lax_numpy_operators_test.py

@@ -255,6 +255,8 @@ def op_record(name, nargs, dtypes, shapes, rng_factory, diff_modes,
               tolerance={dtypes.bfloat16: 4e-2, np.float16: 2e-2,
                          np.float64: 1e-12}),
     op_record("positive", 1, number_dtypes, all_shapes, jtu.rand_default, ["rev"]),
+    op_record("pow", 2, number_dtypes, all_shapes, jtu.rand_positive, ["rev"],
+              tolerance={np.complex128: 1e-14}, check_dtypes=False, alias="power"),
     op_record("power", 2, number_dtypes, all_shapes, jtu.rand_positive, ["rev"],
               tolerance={np.complex128: 1e-14}, check_dtypes=False),
     op_record("rad2deg", 1, float_dtypes, all_shapes, jtu.rand_default, []),
@@ -458,8 +460,8 @@ def testOp(self, op_name, rng_factory, shapes, dtypes, check_dtypes,
     tol = max(jtu.tolerance(dtype, tolerance) for dtype in dtypes)
     if jtu.test_device_matches(["tpu"]) and op_name in (
         "arccosh", "arcsinh", "sinh", "cosh", "tanh", "sin", "cos", "tan",
-        "log", "log1p", "log2", "log10", "exp", "expm1", "exp2", "power",
-        "logaddexp", "logaddexp2", "i0", "acosh", "asinh"):
+        "log", "log1p", "log2", "log10", "exp", "expm1", "exp2", "pow",
+        "power", "logaddexp", "logaddexp2", "i0", "acosh", "asinh"):
       tol = jtu.join_tolerance(tol, 1e-4)
     tol = functools.reduce(jtu.join_tolerance,
                            [tolerance, tol, jtu.default_tolerance()])

tests/lax_numpy_test.py

@@ -5615,7 +5615,7 @@ def testWrappedSignaturesMatch(self):
 
 # TODO(jakevdp): implement missing ufuncs
 UNIMPLEMENTED_UFUNCS = {'spacing', 'bitwise_invert', 'bitwise_left_shift',
-                        'bitwise_right_shift', 'pow', 'vecdot'}
+                        'bitwise_right_shift', 'vecdot'}
 
 
 def _all_numpy_ufuncs() -> Iterator[str]: