XLAShardedTensor.to_local() support

[Hoomaaan]

The implementation adds a to_local() method to XLAShardedTensor class that converts a sharded tensor back to its local representation while preserving gradient information.

1. Core Functionality:

• Returns the global tensor representation containing combined data across all devices
• Maintains the same device placement as the original XLAShardedTensor
• Creates a clone of the global tensor to ensure data independence

2. Gradient Handling:

• Preserves the requires_grad setting from the original tensor
• Maintains gradient values when converting to local representation
• Ensures proper gradient flow through the converted tensor

3. Test Coverage:

• Basic functionality test verifying shape and value preservation
• Dedicated gradient flow test ensuring:
  requires_grad property is preserved
  gradients are properly calculated and maintained
  backward pass works correctly through the local tensor
  gradient values are accurately preserved

[bfolie]

I don't think this method is necessary if there's no additional setup logic

[bfolie]

remove