RFC: add support for negative axes in `permute_dims`

[songyesog2000]

Hi there,

I noticed that negative index can not used for permute_dim when using Jax.numpy namespace.
Although documentation about permute_dims did mention that
"axes (Tuple[int, ...]) – tuple containing a permutation of (0, 1, ..., N-1) where N is the number of axes (dimensions) of x"
However, indexing page says api standard compatible with negative indexing.

Right now, I normalized index before calling permute_dims. However, I do hope your team can make it more clear by whatever measurements. For example, just simply Noting out the non-negative requirement, vs using negative index common people will do, if this is made by designing.

By the way, it will be appreciated if you can share your design philosophy since I am working on a project involves array manipulations. Is that because permute_dim is more like an internal function?

[lucascolley]

just simply Noting out the non-negative requirement,

It may be a good idea to present this requirement more clearly, yes.

it will be appreciated if you can share your design philosophy since I am working on a project involves array manipulations. Is that because permute_dim is more like an internal function?

I would imagine it is because some existing library did not support negative axes (perhaps JAX, as you pointed out already), but maybe @kgryte knows more.

[kgryte]

The current specification stems from its original inclusion as transpose (ref: #228 and #247). The original spec was based on NumPy pre-2.0 (ref: https://numpy.org/doc/1.26/reference/generated/numpy.transpose.html), where it was not explicitly discussed whether to allow negative indices.

That said, personally, I am open to updating the specification to allow negative indices in permute_dims. This is supported in NumPy (ref: https://numpy.org/doc/stable/reference/generated/numpy.permute_dims.html) and would ensure greater consistency in the spec.

[kgryte]

Quick ecosystem survey suggests the following level of support for negative indices:

• Support:
  • NumPy
  • PyTorch (via permute)
  • Dask (via transpose)
• Does not support:
  • JAX
• Unknown
  • CuPy

[rgommers]

Unknown: CuPy

CuPy doesn't yet have permute_dims, but its transpose function does support negative indices:

>>> import cupy as cp
>>> x = cp.arange(24).reshape((2,3,4))
>>> cp.all(cp.transpose(x, axes=(-1,-2,0)) == cp.transpose(x, axes=(2,1,0)))
array(True)

I looked through the JAX issue tracker, and there aren't many mentions of negative axis/axes. jax-ml/jax#10916 (comment) contains the clearest explanation (tl;dr because XLA doesn't support it). However, more recently an allow_negative_indices appeared (see https://docs.jax.dev/en/latest/changelog.html#jax-0-5-3-mar-19-2025), so maybe there's movement there. Unwrapping negative axes at the jax.numpy level or at least somewhere before passing axes to XLA should be feasible I'd think.

@jakevdp do you have any new insights here?

[kgryte]

@jakevdp Following up, do you have thoughts for JAX on whether negative indices can be supported in permute_dims?

[kgryte]

As discussed during a previous workgroup meeting from 2025/09/18, JAX is onboard with supporting, so we are clear to add this to the v2025 specification.