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

[Mosaic TPU] Improve vector layout inference for vector.shape_cast #23792

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

[Mosaic TPU] Improve vector layout inference for vector.shape_cast #23792

wants to merge 1 commit into from

Commits on Sep 20, 2024

  1. [Mosaic TPU] Improve vector layout inference for vector.shape_cast 

    This both simplifies and strengthens the pass. The loop over candidate result implicit
dims was not useful beyond the initial no-op version. I don't think that there's any
implicit dim that would in fact satisfy the checks in a way that the current implementation
doesn't.

What's more important, however, is that the previous implementation never tried to alter
the implicit dim of the input layout, which caused unnecessary failures when inputs came
in with an implicit_dim. Now the rules are more explicit and will proactively adjust the
input layout to make the shape_cast realizable.

PiperOrigin-RevId: 676784915
    @apaszke @Google-ML-Automation
    apaszke authored and Google-ML-Automation committed Sep 20, 2024
    Configuration menu
    Copy the full SHA
    9f31d0e View commit details
    Browse the repository at this point in the history