Merge pull request #19270 from jakevdp:permute-dims-sig

PiperOrigin-RevId: 596992011
jax-ml · Jan 9, 2024 · e76f514 · e76f514
2 parents 2356e57 + 707657e
commit e76f514
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
@@ -543,9 +543,9 @@ def transpose(a: ArrayLike, axes: Sequence[int] | None = None) -> Array:
 
 
 @util._wraps(getattr(np, "permute_dims", None))
-def permute_dims(x: ArrayLike, /, axes: tuple[int, ...]) -> Array:
-  util.check_arraylike("permute_dims", x)
-  return lax.transpose(x, axes)
+def permute_dims(a: ArrayLike, /, axes: tuple[int, ...]) -> Array:
+  util.check_arraylike("permute_dims", a)
+  return lax.transpose(a, axes)
 
 
 @util._wraps(getattr(np, 'matrix_transpose', None))