Skip to content

Conversation

@samanklesaria
Copy link
Collaborator

This PR improves the variable sharding API. Specifically:

  • It renames sharding_names to sharding_metadata (still supporting sharding_names as a property)
  • It generalizes the spmd APIs to allow sharding_metadata to be a NamedSharding

@samanklesaria samanklesaria force-pushed the sharding_metadata branch 4 times, most recently from 5cbafeb to 6611d4b Compare November 17, 2025 16:25
def __getattr__(self, name: str) -> tp.Any:
if name == 'sharding_names': # for backward compatibility
warnings.warn("'sharding_names' is deprecated. Use 'sharding_metadata' instead.", DeprecationWarning)
return self.sharding_metadata
Copy link
Collaborator

Choose a reason for hiding this comment

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

we should use has_metadata and get_metadata in the lines bellow and teach these methods how to handle sharding_names.

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

I note that get_metadata makes a copy of the metadata whenever you call it. This isn't the case when we use the getattr method here normally. Does this matter?

@samanklesaria samanklesaria requested a review from cgarciae January 5, 2026 21:08
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants