[Breaking Change] Implement argsort and make sort numpy API compliant

[Berkant03]

Due Diligence

• General:
  • title of the PR is suitable to appear in the Release Notes
• Implementation:
  • unit tests: all split configurations tested
  • unit tests: multiple dtypes tested
  • NEW unit tests: MPS tested (1 MPI process, 1 GPU)
  • benchmarks: created for new functionality
  • benchmarks: performance improved or maintained
  • documentation updated where needed

Description

argsort and sort have been changed to reflect the Numpy API, while also allowing old usage.

Issue/s resolved: #777

Type of change

• New feature

Does this change modify the behaviour of other functions? If so, which?

yes, sort now does not give back the indices when sorting like pytorch. Instead a flag must be given to achieve same behaviour.

[brownbaerchen]

Thanks! I think it's a good idea to be closer to Numpy and to support argsort. However, I don't think we should silently ignore additional arguments that Numpy users are used to.

Note that the errors in the collection step in pytest happen all over the place, including in other repositories, and have nothing to do with this PR.

[brownbaerchen]

Excellent! If you apply something similar to argsort, this should be ready to merge.

[brownbaerchen]

Thanks! I think we can clean up the tests a bit more. Otherwise, this is ready in my opinion.

[Berkant03]

The current tests fail, due to an API change in numpy 2.5, which added the descending keyword to argsort, but is absent in previous versions. How should we handle the cases with a lower numpy version? @brownbaerchen

[brownbaerchen]

The current tests fail, due to an API change in numpy 2.5, which added the descending keyword to argsort, but is absent in previous versions. How should we handle the cases with a lower numpy version? @brownbaerchen

That's annoying! I think the best we can do is to split the tests into ones with descending and without. Then you can skip tests using descending if np.lib.NumpyVersion(np.__version__) < '2.5.0'.
You can use a decorator:

class Tests(...):
    @unittest.skipif(np.lib.NumpyVersion(np.__version__) < '2.5.0', f'Numpy version {np.__version__} does not support `descending` keyword in `sort`)
    def test_sort(self):
        ...
    def test_sort_no_descending_argument(self):
        ...

Alternatively, you can use the unittest.subtest stuff within a single function and skip some subtests from there.
Neither is particularly clean.. Do use skipping infrastructure, though. This will tell you how many tests were skipped during the collection step and why and is not silent.

[codecov]

Codecov Report

❌ Patch coverage is 50.00000% with 12 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.74%. Comparing base (c26de1f) to head (e3072e3).
⚠️ Report is 6 commits behind head on main.
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
heat/core/manipulations.py	45.45%	12 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2386      +/-   ##
==========================================
- Coverage   91.81%   91.74%   -0.08%     
==========================================
  Files          87       87              
  Lines       14149    14168      +19     
==========================================
+ Hits        12991    12998       +7     
- Misses       1158     1170      +12     

Flag	Coverage Δ
unit	91.74% <50.00%> (-0.08%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

[brownbaerchen]

Sorry, I have a few more small points, but should be quick to address 🙈

[brownbaerchen]

Feel free to remove any unnecessary tests, or not. Otherwise, we can finally merge :D

[brownbaerchen]

The failures in previous pushes look unrelated to this.. No idea what's going on there... :(

[Berkant03]

@brownbaerchen, I removed the specific test. The codebase test fails due to not enough storage. If everything is fine, you can merge it.

[brownbaerchen]

@brownbaerchen, I removed the specific test. The codebase test fails due to not enough storage. If everything is fine, you can merge it.

Let's merge! @JuanPedroGHM, do you have some bandwidth for looking into why we are suddenly running out of disk space on codebase during installation?