Distributed `vmap` functionality for vectorization across split dimension

[mrfh92]

Due Diligence

• General:
  • [x] title of the PR is suitable to appear in the Release Notes
• Implementation:
  • [x] unit tests: all split configurations tested
  • [x] unit tests: multiple dtypes tested
  • [x] documentation updated where needed

Description

Issue/s resolved: #1383 The implemented vmap-functionality is similar to the one of PyTorch. It allows vectorized application of PyTorch-functions to one or multiple DNDarrays with vectorization taking place along the respective split axes.

An artificial example, how vmap does work, is found in the tests.

Type of change

new feature

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

no

[github-actions[bot]]

Thank you for the PR!

[github-actions[bot]]

Thank you for the PR!

[github-actions[bot]]

Thank you for the PR!

[github-actions[bot]]

Thank you for the PR!

[github-actions[bot]]

Thank you for the PR!

[github-actions[bot]]

Thank you for the PR!

[github-actions[bot]]

Thank you for the PR!

[github-actions[bot]]

Thank you for the PR!

[github-actions[bot]]

Thank you for the PR!

[codecov[bot]]

Codecov Report

Attention: Patch coverage is 94.28571% with 2 lines in your changes missing coverage. Please review.

Project coverage is 91.91%. Comparing base (a774559) to head (b388446). Report is 2 commits behind head on main.

Files	Patch %	Lines
heat/core/vmap.py	94.11%	2 Missing :warning:

Additional details and impacted files

```diff @@ Coverage Diff @@ ## main #1516 +/- ## ======================================= Coverage 91.91% 91.91% ======================================= Files 80 81 +1 Lines 11942 11977 +35 ======================================= + Hits 10976 11009 +33 - Misses 966 968 +2 ``` | [Flag](https://app.codecov.io/gh/helmholtz-analytics/heat/pull/1516/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=helmholtz-analytics) | Coverage Δ | | |---|---|---| | [unit](https://app.codecov.io/gh/helmholtz-analytics/heat/pull/1516/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=helmholtz-analytics) | `91.91% <94.28%> (+<0.01%)` | :arrow_up: | Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=helmholtz-analytics#carryforward-flags-in-the-pull-request-comment) to find out more.

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

[github-actions[bot]]

Thank you for the PR!

[mrfh92]

@ClaudiaComito problem should be fixed now. In fact, we require PyTorch >= 2.0.0 for vmap as vmap is only available in these versions.

[github-actions[bot]]

Thank you for the PR!

[mrfh92]

@ClaudiaComito codecov was a bit lower than expected by adding the uncovered line (~-3%), so I improved coverage

[github-actions[bot]]

Thank you for the PR!

[github-actions[bot]]

Thank you for the PR!

[mrfh92]

@ClaudiaComito @JuanPedroGHM thx for reviewing :+1: