Optimal Device Assignment for 2-D mesh

[Lumosis]

Description

Introduces an algorithm to derive the optimal device assignment that maps a 1D ring of logical devices to a 2D mesh of physical devices.

Tests

* Baseline with optimize_mesh_for_tpu_v6e=false

================================================
MLPerf Results Summary
================================================
SUT name : PySUT
Scenario : Offline
Mode     : PerformanceOnly
Samples per second: 17.2537
Tokens per second: 4869.9130
Result is : VALID
  Min duration satisfied : Yes
  Min queries satisfied : Yes
  Early stopping satisfied: Yes



* Baseline with optimize_mesh_for_tpu_v6e=true

================================================
MLPerf Results Summary
================================================
SUT name : PySUT
Scenario : Offline
Mode     : PerformanceOnly
Samples per second: 19.0753
Tokens per second: 5414.4059
Result is : VALID
  Min duration satisfied : Yes
  Min queries satisfied : Yes
  Early stopping satisfied: Yes


* New code

================================================
MLPerf Results Summary
================================================
SUT name : PySUT
Scenario : Offline
Mode     : PerformanceOnly
Samples per second: 19.5143
Tokens per second: 5528.4437
Result is : VALID
  Min duration satisfied : Yes
  Min queries satisfied : Yes
  Early stopping satisfied: Yes

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• [x] I have performed a self-review of my code.
• [x] I have necessary comments in my code, particularly in hard-to-understand areas.
• [x] I have run end-to-end tests tests and provided workload links above if applicable.
• [x] I have made or will make corresponding changes to the doc if needed.

[github-actions[bot]]

This PR has been automatically marked as stale because it has not had recent activity. It will be closed soon if no further activity occurs. Thank you for your contributions.

[github-actions[bot]]

This PR was closed because it has been inactive for a while. Please reopen it if you are still working on it.