Could you add a test for a large matrix, I am slightly worried about this only accidentally working. The problem is that if the input is exactly identical to the output, we probably pass this on into lapack (einsum?) directly.
Even after a quick check, I am not sure that the matmul inner-loop is actually able to deal with overlap just fine.

Could you add a test for a large matrix, I am slightly worried about this only accidentally working.

I just added a dedicated test case in 1e62199; even when upping the size of the matrix from 10**6 to 10**8 I'm still getting the expected result. I assume these results are somewhat promising?

Thanks, I like to have the test in either case, I will have a closer look at the code.

I have to check if and when the iterator makes a copy anyway. If it does (which is plausible) writing a @= b vs. a = a @ b has no advantage (in fact it is probably worse), but it will be safe. If it does not copy, things get a bit tricky, because matmul has a bunch of different code paths that may be taken, and we have to check/test all of them.

Sorry, my bad, I missed/forgot that we of course already have a matmul special case that always forces the copy. So this is safe, but slow/not very useful.

Especially cases where a broadcast could happen, if a is (n, n) and b is (n, 1).

Great point, and potentially tricky :/. I remember fixing some things around this and making it more strict, but we did not disallow the output to broadcast the inputs at that time. This means that this is still incorrect:

In [4]: arr = np.array([1., 2.])

In [5]: np.matmul(arr, arr, out=arr)
Out[5]: array([5., 5.])

There could be an argument to allow the output to broadcast in some cases, but when inplace ops are used for a gufunc we are on very thin ice.

EDIT: Since it may not be clear, Eric's example is fine, because it happens in the core dimensions, but we currently do allow the output to broadcast in the non-core dimensions. In the above example, the input has more core dimensions than the output, so an input core dimension because becomes an outer output dimension.

Can you add a test for what happens when a @ b and a are not the same shape? Especially cases where a broadcast could happen, if a is (n, n) and b is (n, 1).

Currently it raises an exception, rather than automatically broadcasting the output. In this sense it's somewhat unique compared to the other inplace operation, but then again __matmul__ is also the only dunder that can produce an array that is smaller in size compared to the original.

All in all I think I might prefer the current (more explicit) ValueError over an implicit broadcast (this sounds like something that could easily bite someone in the ass).

In [1]: a = np.arange(9).reshape(3, 3)
   ...: b = np.ones((3, 1), dtype=int)

In [2]: a @= b
---------------------------------------------------------------------------
ValueError: matmul: Output operand 0 has a mismatch in its core dimension 1, with gufunc signature (n?,k),(k,m?)->(n?,m?) (size 3 is different from 1)

On a unrelated note: either sphinx or one its dependencies is causing issues in circleci.
Based on the exception a package requiring python >= 3.9 is being installed on python 3.8.

Traceback (most recent call last):
  File "/home/circleci/repo/venv/lib/python3.8/site-packages/sphinx/cmd/build.py", line 280, in build_main
    app.build(args.force_all, filenames)
  File "/home/circleci/repo/venv/lib/python3.8/site-packages/sphinx/application.py", line 343, in build
    self.builder.build_update()
  File "/home/circleci/repo/venv/lib/python3.8/site-packages/sphinx/builders/__init__.py", line 293, in build_update
    self.build(to_build,
  File "/home/circleci/repo/venv/lib/python3.8/site-packages/sphinx/builders/__init__.py", line 307, in build
    updated_docnames = set(self.read())
  File "/home/circleci/repo/venv/lib/python3.8/site-packages/sphinx/builders/__init__.py", line 412, in read
    self._read_parallel(docnames, nproc=self.app.parallel)
  File "/home/circleci/repo/venv/lib/python3.8/site-packages/sphinx/builders/__init__.py", line 463, in _read_parallel
    tasks.join()
  File "/home/circleci/repo/venv/lib/python3.8/site-packages/sphinx/util/parallel.py", line 108, in join
    if not self._join_one():
  File "/home/circleci/repo/venv/lib/python3.8/site-packages/sphinx/util/parallel.py", line 129, in _join_one
    raise SphinxParallelError(*result)
sphinx.errors.SphinxParallelError: AttributeError: 'PosixPath' object has no attribute 'readlink'

Marking for triage-review for now. I am not quite sure what to do about my example where the number of core dimensions of input and output differs; can we ignore it?
A possible work-around to reject it (without modifying the general case) would be to pass axes=[(-1,), None, (-1,)] to ensure that input and output have the same number of core dimensions (or (-2, -1) of course)).
However, it seems like it is not possible to pass None for the second input and I am not sure how hard it would be to add that functionality (it feels like it shouldn't be hard, but not sure).

I am not quite sure what to do about my example where the number of core dimensions of input and output differs; can we ignore it?

I would be in favor of ignoring it and just letting it raise (i.e. the current behavior). There a few other functions that reduce either the dimensionality or alter the output shape w.r.t. the input shape (e.g. np.mean) and those raise a ValueError as well, even though the result could in principle be broadcasted.

In [1]: a = np.arange(9).reshape(3, 3)
   ...: b = np.ones((3, 1), dtype=int)

In [2]: b.mean(out=a)
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
  ...
ValueError: output parameter for reduction operation add has the wrong number of dimensions: Found 2 but expected 0

@BvB93 sorry, I should have repeated my example:

In [1]: arr = np.array([1.0, 2.0])

In [2]: arr @ arr
Out[2]: 5.0

In [3]: arr @= arr

In [4]: arr
Out[4]: array([5., 5.])

close/reopen to regenerate CI logs

I pulled the nuclear option on this, because the tests noticed the array coercion. This means that the errors are pretty terrible, maybe most impressive:

In [2]: arr = np.ones((3, 3))
In [3]: arr @= np.ones(3)
ValueError: axes item 1 should be a tuple with a single element, or an integer

but at least the ValueError part is technically correct :cough:.

By "nuclear" you mean call back out to python with np.matmul(a, b, out=a, axes=[(-2, -1), (-2, -1), (-2, -1)]). How do you think we could get a better error message?

Yeah, that is what I mean, the initial code converted the b to an array in a @= b, which makes it easy to check things up-front.

We can improve the generic message a little, but that will still be a bit confusing probably. I suspect getting a nice one would actually require to manually replace it here (maybe pragmatically assuming that the other side isn't a super weird object).

Something like "matmul operand 2 has 1 dimension, but 2 dimensions are required (axes tuple has length 2)." (EDIT: add the ufunc name, we should give it.)

OK, now based on changing the other error to use AxisError, replacing it here now (a bit of a best effort since you could do evil things with nested object arrays). That removes the worst of the errors (which referred to the axes argument the user never even supplied).

I think we should put this in and work out better error messages as they arise. For the case above, the message is now:

>>> arr = np.ones((3, 3))
>>> arr @= np.ones(3)
Traceback (most recent call last):
  File "<console>", line 1, in <module>
ValueError: inplace matrix multiplication requires the first operand to have \
    at least one and the second at least two dimensions.

A casting error also seems clear enough:

>>> arr = np.ones((3, 3), dtype=int)
>>> arr @= np.ones((3,3))
Traceback (most recent call last):
  File "<console>", line 1, in <module>
numpy.core._exceptions._UFuncOutputCastingError: Cannot cast ufunc 'matmul' output \
    from dtype('float64') to dtype('int64') with casting rule 'same_kind'

I finished this from my side, if you are happy with the new solution.

Thanks @BvB93 and @seberg.