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

#sdy Handle redundant shard maps during import. #16943

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

#sdy Handle redundant shard maps during import. #16943

wants to merge 1 commit into from
+129 −18

Commits on Sep 9, 2024

  1. #sdy Handle redundant shard maps during import. 

    Even with a mesh with multiple axes size > 1, if the shmap is operating on no axes > 1, then JAX will create the call but not have the pattern of the custom calls. So just inline.

PiperOrigin-RevId: 672536491
    @bartchr808 @copybara-github
    bartchr808 authored and copybara-github committed Sep 9, 2024
    Configuration menu
    Copy the full SHA
    d377150 View commit details
    Browse the repository at this point in the history