Add support for matrix transpose

[aaishwarymishra]

• ☑️ I understand it is strictly prohibited to use AI to write PR description
  Added matrix transpose .mT property for the Array API compliance. Part of MLX compatibility: Constants, Datatypes and Array Attributes data-apis/array-api-compat#451
• AI usage disclosure:
  None

[nastya236]

Thanks for your contribution!
However, I’m not sure we want to support this kind of behaviour. Once a tensor has more than two dimensions, there is no universally obvious “transpose.” Therefore, requiring the permutation to be explicit potentially avoids confusion and makes code more readable.

Also, in LLM workloads, a typical permutation is something like: queries = queries.transpose(0, 2, 1, 3) before applying RoPE in the attention layer. Therefore, I don’t see a reason to treat transposing the last two axes as a special case while not doing the same for other common permutations..
Is there a specific use case you have in mind where swapping the last two axes via x.mT is particularly useful or common?

[aaishwarymishra]

Hi, the main goal of most of my PR's is to make MLX Array API compliant, https://data-apis.org/array-api/latest/index.html just like numpy and jax so the code can be made backend agnostic downstream :)

[ev-br]

Once a tensor has more than two dimensions, there is no universally obvious “transpose.”

Exactly! Which is precisely the reason NumPy arrays have two attributes: .T and .mT. The former is (backwards compatible) revert all axes, and the latter only transposes the two trailing dimensions. To illustrate:

In [1]: import numpy as np

In [4]: a = np.arange(2*3*4*5).reshape(2, 3, 4, 5)

In [6]: a.shape
Out[6]: (2, 3, 4, 5)

In [7]: a.T.shape
Out[7]: (5, 4, 3, 2)

In [8]: a.mT.shape
Out[8]: (2, 3, 5, 4)

The interpretation of .mT ("matrix transpose") is that an array represents a batch of 2D matrices, where the two last axes are the matrix dimensions, and leading dimensions are batch dimensions. In the example above, a of shape (2, 3, 4, 5) represents a (2, 3)-shaped batch of (4, 5)-shaped matrices.

Is there a specific use case you have in mind where swapping the last two axes via x.mT is particularly useful or common?

Absolutely! One typical use case is linear algebra with batched matrices. For example, a normal equation solution for the linear regression problem is just (modulo inv vs solve, of course)

inv( X.mT @ X) @ X.mT @ y

Note how this formula "just works" for solving multiple regressions at once, for a batched design matrix and/or batched y.

---

Most Array API compatible array libraries implement the .mT attribute with numpy-compatible behavior:

In [9]: jnp.asarray(a).mT.shape       # jax.numpy
Out[9]: (2, 3, 5, 4)

In [10]: jnp.asarray(a).T.shape
Out[10]: (5, 4, 3, 2)

In [11]: torch.asarray(a).mT.shape     # pytorch
Out[11]: torch.Size([2, 3, 5, 4])

Finally, while the Array API spec itself requires that the "usual" .T attribute only works for strictly 2D arrays, but AFAIU the majority of real array libraries deviate from that, and match NumPy instead.

All in all, and as far as experience with Array API is any guide, it is best that an array library supports both .T and .mT, where the latter only tranposes two last dimensions.

[nastya236]

Regarding being framework agnostic, we don't want to mirror every API or behaviour from other frameworks that will unnecessarily increase the API and codebase, but rather focus on what makes sense for MLX and its use cases.

Regarding mT, personally, I would prefer to explicitly specify the axes in transpose, rather than adding a special-case API for a particular permutation (even though I agree that this is a valid batched matrix multiplication use case).

That being said, I'm not opposed to adding it. My main concern is that I don't think “being framework-agnostic” is a strong enough reason for introducing the API. I think the rationale should instead be that x.mT provides a useful and common abstraction which I don't believe it is, but maybe I am biased..

May I ask what do you think, @zcbenz ?

[ev-br]

Re framework agnosticism I should probably explain our motivation a bit better. The short story is that Array API defines a minimum common API surface for array libraries. The primary use case is a higher-level library (scipy, scikit-learn, ...) using an Array API producer (pytorch, jax, cupy, dpnp, hopefully MLX) instead of just numpy.

The way it works is that the internals of the consumer library are implemented against the Array API---and then users can just switch the array library: they feed MLX arrays to SciPy functions, all array processing is done by MLX, and users get MLX arrays back.
This is not hypothetical at all: SciPy is planning to make Array API support public in the next release, and scikit-learn will likely follow suit.

Back to .mT: indeed, the last two axes are not that special. However any SciPy or scikit-learn function which expects .mT to exist (meaning: quite a few!) will be blocked, and users will have no way to fix it---just because the failing call to .mT happens inside scikit-learn!

Of course, the full compatibility is not achievable, and is not even a goal. We fully understand that MLX focuses on what makes sense for MLX! Consumer libraries will work around fundamental design constraints (complex128, data-dependent shapes and so on). Small differences however (.mT being an example) seem to have a more favorable cost/benefit ratio: assuming it does not cost excessively much for MLX, the benefit for downstream consumers is quite large.

The wall of text above is just hopefully useful as a high-level motivation of the stream of Array API PRs, and hopefully provides a use case for MLX maintainers to consider.
Am happy to go into details if helpful, or leave things as they are if not :-).

EDIT: Here's the canonical "purpose and scope" text for the Array API movement, https://data-apis.org/array-api/latest/purpose_and_scope.html

[zcbenz]

@nastya236 I'm generally open to conforming to Array API spec when it adds no maintenance burden, for example adding an acos alias for arccos API, mT is a lot more than that, but they do make a point about it being useful to scipy/scikit-learn users and it would be pretty cool if MLX can be a new backend for them. I also found this relevant issue from Awni when searching: scikit-learn/scikit-learn#29673.

So I think I'm not fond of the API itself, but if it is blocking scipy/scikit-learn from adopting MLX as a new backend, we should accept it.

On the PR itself, since this is not a trivial one-liner, and there is also a functional version matrix_transpose, I think we should make it a formal C++ matrix_transpose API and expose it in python bindings.

[ev-br]

I also found this relevant issue from Awni when searching: scikit-learn/scikit-learn#29673.

Just to add that in SciPy there is no dedicated issue for MLX specifically. There's however a body of "lore", spread over multiple issues and PR reviews over the last few years we've been working on the Array API support, where it's clear that there's significant interest for supporting "metal backends".
Specific work did not materialize yet, with a significant hurdle being precisely the large body of small inconsistencies with the spec.