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

[rpsilva-aws]

The grad doesn't need to be cloned? Is it fine that we break the reference for the tensor, but not its grad? See if we can add a test that shows it is the case (e.g. both the prior and newer tensor updating the same grad).

[Hoomaaan]

Updated to clone the self.grad if available

[rpsilva-aws]

nit: result.requires_grad_(self.requires_grad) for in-place.

[Hoomaaan]

Done!