Skip to content

fix all_gather_into_tensor test and logic#9332

Merged
bfolie merged 2 commits intomasterfrom
bfolie/all-gather-into-tensor
Jun 16, 2025
Merged

fix all_gather_into_tensor test and logic#9332
bfolie merged 2 commits intomasterfrom
bfolie/all-gather-into-tensor

Conversation

@bfolie
Copy link
Copy Markdown
Collaborator

@bfolie bfolie commented Jun 10, 2025

Most cases now work, except for stacking + compilation, which requires a change to upstream pytorch (pytorch/pytorch#155632)

@bfolie bfolie requested review from bhavya01, pgmoka and qihqi June 10, 2025 23:10
Comment thread torch_xla/distributed/xla_backend.py
Comment thread test/pjrt/test_collective_ops_tpu.py
Comment thread torch_xla/distributed/xla_backend.py Outdated
Comment thread torch_xla/distributed/xla_backend.py
@bfolie bfolie requested review from bhavya01 and pgmoka June 16, 2025 16:19
Copy link
Copy Markdown
Collaborator

@pgmoka pgmoka left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM minus some very small nits and suggestions.

Comment thread test/pjrt/test_collective_ops_tpu.py
Comment thread test/pjrt/test_collective_ops_tpu.py
@bfolie bfolie merged commit 9b0b02f into master Jun 16, 2025
40 of 42 checks passed
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.

3 participants