Fix for double counting error from #2093

[brownbaerchen]

This PR fixes one of the issues in #2093, namely the bug with running code:

import heat as ht

split = 0
shape = (4, 3)

A = ht.ones(shape, split=split)
B = ht.ones(shape[::-1], split=split)

C = A @ B
C_glob = (A.resplit(None) @ B.resplit(None)).resplit(C.split)

print(C)
print(C_glob)

which gives an output when run with two tasks:

DNDarray([[3., 3., 3., 3.],
          [6., 6., 6., 6.],
          [3., 3., 3., 3.],
          [6., 6., 6., 6.]], dtype=ht.float32, device=cpu:0, split=0)
DNDarray([[3., 3., 3., 3.],
          [3., 3., 3., 3.],
          [3., 3., 3., 3.],
          [3., 3., 3., 3.]], dtype=ht.float32, device=cpu:0, split=0)

This PR implements an analogous test with four tasks. I first pushed the test without the fix to record that the test fails on the current main. Then, I pushed the fix, which is just two small typos.

This does not close the original issue because it doesn't resolve all the issues, but I thought this is something we can merge quickly and independently of a larger refactor and fix for the other problems, which will come later.

[Copilot]

Pull request overview

This PR addresses the double-counting edge case reported in #2093 by adding a targeted regression test and adjusting matmul()’s remainder-flag logic to avoid incorrect handling when block sizes are 1.

Changes:

• Added an MPI-size-specific regression test for a matmul edge case from #2093.
• Fixed matmul() remainder checks to use the correct outer-dimension block sizes (mB/nB) instead of kB.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
tests/core/linalg/test_basics.py	Adds a regression test covering the #2093 edge case for distributed matmul.
heat/core/linalg/basics.py	Corrects remainder-flag conditions in matmul() to prevent incorrect behavior/double counting in edge cases.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.