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 2 times, most recently from a734ef9 to d8df0d7 Compare November 13, 2025 23:16
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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants