XLAShardedTensor.to_local() support - #9505

Open

Hoomaaan wants to merge 13 commits into

pytorch:masterfrom

Hoomaaan:toLocal_wspec

Commits on Sep 5, 2025

:qImplement XLAShardedTensor._spec and test

aws-cph
authored and
Hoomaaan
committed
ed86274
Removed auto wrapping sharding propagation, added cached spec invalidation

aws-cph
authored and
Hoomaaan
committed
858863e
Removing lazy import

aws-cph
authored and
Hoomaaan
committed
dd79690
Added test for catching thrown error in spec

aws-cph
authored and
Hoomaaan
committed
8c43478
Test for Routing XLA device handling through distribute_tensor to ensure proper XLA support and maintain consistency with PyTorch/XLA SPMD integration.
Hoomaaan
committed
4b5af4b
[XLA] Implement XLAShardedTensor.to_local()
Hoomaaan
committed
df9279b
run git_fix for yapf
Hoomaaan
committed
bf331da
Remove print statement
Hoomaaan
committed
04c5cd3
code clean up
Hoomaaan
committed
52a5e70
Remove redundant setUpClass constructor
Hoomaaan
committed
4c7ffc2
Clone the grads and use inplace method for requires_grad
Hoomaaan
committed
bd9c9f3
fix the failing CI by reverting to default requires_grad
Hoomaaan
committed
be7ab62
run yapf
Hoomaaan
committed
b378a46