Skip to content

Fix PartitionSpec hardcoding in resharding_engine for multi-dimensional TPU meshes.#68

Open
copybara-service[bot] wants to merge 1 commit into
mainfrom
test_927563534
Open

Fix PartitionSpec hardcoding in resharding_engine for multi-dimensional TPU meshes.#68
copybara-service[bot] wants to merge 1 commit into
mainfrom
test_927563534

Conversation

@copybara-service
Copy link
Copy Markdown

@copybara-service copybara-service Bot commented Jun 6, 2026

Fix PartitionSpec hardcoding in resharding_engine for multi-dimensional TPU meshes.

This change ensures that multihost_utils.global_array_to_host_local_array dynamically uses the sharding spec of the input array, improving support for various TPU mesh topologies.

@google-cla
Copy link
Copy Markdown

google-cla Bot commented Jun 6, 2026

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

…al TPU meshes.

This change ensures that `multihost_utils.global_array_to_host_local_array` dynamically uses the sharding spec of the input array, improving support for various TPU mesh topologies.

PiperOrigin-RevId: 927563534
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

0 participants