Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[sharding_in_types] If out_aval.sharding is not None and the user specified out_sharding is None, concretize it with the device assignment available and add it to the final out_shardings that's used for lowering and compilation. #24387

Merged
merged 1 commit into from
Oct 18, 2024
Merged

[sharding_in_types] If out_aval.sharding is not None and the user specified out_sharding is None, concretize it with the device assignment available and add it to the final out_shardings that's used for lowering and compilation. #24387

merged 1 commit into from
Oct 18, 2024
+43 −26

Commits on Oct 18, 2024

  1. [sharding_in_types] If out_aval.sharding is not None and the user spe… 

    …cified out_sharding is None, concretize it with the device assignment available and add it to the final out_shardings that's used for lowering and compilation.

This will allow us to return the exact sharding spec that sharding propagation rules figured out.

PiperOrigin-RevId: 687349015
    @yashk2810 @Google-ML-Automation
    yashk2810 authored and Google-ML-Automation committed Oct 18, 2024
    Configuration menu
    Copy the full SHA
    2153de4 View commit details
    Browse the repository at this point in the history