PR #9505 XLAShardedTensor.to_local() support

XLAShardedTensor.to_local() support #9505

Hoomaaan wants to merge 13 commits into pytorch:master from Hoomaaan:toLocal_wspec

Hoomaaan force pushed from aa0c9557 to f0c89b94 265 days ago

Hoomaaan force pushed from f0c89b94 to 933a964e 265 days ago

rpsilva-aws requested a review from

bfolie 259 days ago

rpsilva-aws requested a review from

zhanyong-wan 259 days ago

rpsilva-aws requested a review from

rpsilva-aws 259 days ago

bfolie commented on 2025-08-02

Hoomaaan force pushed from 812a69aa to 6dc23517 246 days ago

Hoomaaan requested a review from

bfolie 246 days ago

bfolie approved these changes on 2025-08-14

rpsilva-aws requested changes on 2025-08-15

rpsilva-aws approved these changes on 2025-08-15

:qImplement XLAShardedTensor._spec and test

ed862748

Removed auto wrapping sharding propagation, added cached spec invalid…

858863e2

Removing lazy import

dd79690e

Added test for catching thrown error in spec

8c434788

Test for Routing XLA device handling through distribute_tensor to ens…

4b5af4bf

[XLA] Implement XLAShardedTensor.to_local()

df9279b1

run git_fix for yapf

bf331da5

Remove print statement

04c5cd36

code clean up

52a5e705

Remove redundant setUpClass constructor

4c7ffc28

Clone the grads and use inplace method for requires_grad

bd9c9f37

fix the failing CI by reverting to default requires_grad

be7ab62e

run yapf

b378a462

Hoomaaan force pushed from 3aaae842 to b378a462 222 days ago

qihqi enabled auto-merge (squash) 217 days ago

Reviewers

rpsilva-aws

bfolie

zhanyong-wan

Assignees

No one assigned

Labels

None yet

Milestone

No milestone

xla XLAShardedTensor.to_local() support #9505 Open

XLAShardedTensor.to_local() support #9505

xla
XLAShardedTensor.to_local() support
#9505

Open