Skip to content

Implement XLAShardedTensor._spec and test#9488

Merged
jeffhataws merged 5 commits intopytorch:masterfrom
aws-cph:aws-cph_dtensor_spec
Jul 24, 2025
Merged

Implement XLAShardedTensor._spec and test#9488
jeffhataws merged 5 commits intopytorch:masterfrom
aws-cph:aws-cph_dtensor_spec

Conversation

@aws-cph
Copy link
Copy Markdown
Contributor

@aws-cph aws-cph commented Jul 17, 2025

Implementing and adding tests for XLAShardedTensor._spec in regards to #9418.

@aws-cph aws-cph force-pushed the aws-cph_dtensor_spec branch from bb4eb3b to 28c4f6f Compare July 17, 2025 20:40
jeffhataws
jeffhataws reviewed Jul 17, 2025
View reviewed changes
Comment thread torch_xla/distributed/spmd/xla_sharded_tensor.py Outdated
jeffhataws
jeffhataws reviewed Jul 17, 2025
View reviewed changes
Comment thread torch_xla/distributed/spmd/xla_sharded_tensor.py Outdated
jeffhataws
jeffhataws reviewed Jul 17, 2025
View reviewed changes
Comment thread torch_xla/distributed/spmd/xla_sharding.py
fhaolinaws
fhaolinaws reviewed Jul 18, 2025
View reviewed changes
Comment thread torch_xla/distributed/spmd/xla_sharded_tensor.py Outdated
fhaolinaws
fhaolinaws reviewed Jul 18, 2025
View reviewed changes
Comment thread torch_xla/distributed/spmd/xla_sharded_tensor.py Outdated
fhaolinaws
fhaolinaws reviewed Jul 18, 2025
View reviewed changes
Comment thread torch_xla/distributed/spmd/xla_sharded_tensor.py
fhaolinaws
fhaolinaws reviewed Jul 18, 2025
View reviewed changes
Comment thread torch_xla/distributed/spmd/xla_sharded_tensor.py
fhaolinaws
fhaolinaws reviewed Jul 18, 2025
View reviewed changes
Comment thread torch_xla/distributed/spmd/xla_sharded_tensor.py
fhaolinaws
fhaolinaws reviewed Jul 18, 2025
View reviewed changes
Comment thread test/spmd/test_xla_dtensor_spec_conversion.py
@aws-cph aws-cph force-pushed the aws-cph_dtensor_spec branch 2 times, most recently from 78a5639 to 4dd06ad Compare July 18, 2025 22:08
bfolie
bfolie reviewed Jul 21, 2025
View reviewed changes
Comment thread torch_xla/distributed/spmd/xla_sharded_tensor.py
bfolie
bfolie reviewed Jul 21, 2025
View reviewed changes
Comment thread torch_xla/distributed/spmd/xla_sharded_tensor.py Outdated
bfolie
bfolie reviewed Jul 21, 2025
View reviewed changes
Comment thread torch_xla/distributed/spmd/xla_sharded_tensor.py Outdated
bfolie
bfolie reviewed Jul 21, 2025
View reviewed changes
Comment thread test/spmd/test_xla_dtensor_spec_conversion.py Outdated
bfolie
bfolie reviewed Jul 21, 2025
View reviewed changes
Comment thread torch_xla/distributed/spmd/xla_sharding.py
qihqi
qihqi reviewed Jul 21, 2025
View reviewed changes
Comment thread torch_xla/distributed/spmd/xla_sharded_tensor.py Outdated
Comment thread torch_xla/distributed/spmd/xla_sharded_tensor.py
Comment thread torch_xla/distributed/spmd/xla_sharded_tensor.py Outdated
@aws-cph aws-cph force-pushed the aws-cph_dtensor_spec branch from 3ab7296 to 16f02a4 Compare July 22, 2025 18:43
@aws-cph aws-cph force-pushed the aws-cph_dtensor_spec branch from 1d2aa97 to 213efc6 Compare July 23, 2025 18:21
bfolie
bfolie reviewed Jul 23, 2025
View reviewed changes
Comment thread test/spmd/test_xla_sharding.py Outdated
@aws-cph aws-cph force-pushed the aws-cph_dtensor_spec branch from 213efc6 to 25147a2 Compare July 23, 2025 21:48
@aws-cph aws-cph force-pushed the aws-cph_dtensor_spec branch from 25147a2 to 5a61b1d Compare July 23, 2025 21:49
@aws-cph aws-cph requested a review from bfolie July 23, 2025 21:49
bfolie
bfolie reviewed Jul 23, 2025
View reviewed changes
Comment thread test/spmd/test_xla_dtensor_spec_conversion.py
bfolie
bfolie approved these changes Jul 23, 2025
View reviewed changes
Copy link
Copy Markdown
Collaborator

@bfolie bfolie left a comment

Choose a reason for hiding this comment

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

Looks good overall. I have one suggestion about test structure but it's non-blocking

@jeffhataws jeffhataws merged commit ca47198 into pytorch:master Jul 24, 2025
23 of 24 checks passed
@aws-cph aws-cph deleted the aws-cph_dtensor_spec branch July 24, 2025 17:40
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@bfolie bfolie bfolie approved these changes

@jeffhataws jeffhataws jeffhataws approved these changes

@rpsilva-aws rpsilva-aws Awaiting requested review from rpsilva-aws

@yaoshiang yaoshiang Awaiting requested review from yaoshiang

@qihqi qihqi Awaiting requested review from qihqi

@fhaolinaws fhaolinaws Awaiting requested review from fhaolinaws

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

5 participants

@aws-cph @qihqi @bfolie @jeffhataws @fhaolinaws