[Pallas] Cast bool masks to float before expanding in _mask_to codegen#2216
Merged
Conversation
AmesingFlank
force-pushed
the
AmesingFlank/stack/37
branch
from
938bde3 to
a1867d1
Compare
AmesingFlank
force-pushed
the
AmesingFlank/stack/38
branch
from
e14cb44 to
e823bb4
Compare
This was referenced May 3, 2026
Merged
meta-cla
Bot
added
the
CLA Signed
AmesingFlank
added a commit
that referenced
this pull request
May 3, 2026
Mosaic TPU compiler cannot reshape bool vectors (vector<Nxi1> -> vector<Nx1xi1>). When _mask_to generates combined masks like mask_q[:, None] & mask_kv[None, :], each mask is now cast to float32 before dimension expansion and combined with * instead of &. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> stack-info: PR: #2216, branch: AmesingFlank/stack/38
AmesingFlank
force-pushed
the
AmesingFlank/stack/38
branch
from
e823bb4 to
8e224a5
Compare
AmesingFlank
added a commit
that referenced
this pull request
May 4, 2026
Mosaic TPU compiler cannot reshape bool vectors (vector<Nxi1> -> vector<Nx1xi1>). When _mask_to generates combined masks like mask_q[:, None] & mask_kv[None, :], each mask is now cast to float32 before dimension expansion and combined with * instead of &. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> stack-info: PR: #2216, branch: AmesingFlank/stack/38
AmesingFlank
added a commit
that referenced
this pull request
May 4, 2026
Mosaic TPU compiler cannot reshape bool vectors (vector<Nxi1> -> vector<Nx1xi1>). When _mask_to generates combined masks like mask_q[:, None] & mask_kv[None, :], each mask is now cast to float32 before dimension expansion and combined with * instead of &. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> stack-info: PR: #2216, branch: AmesingFlank/stack/38
AmesingFlank
added a commit
that referenced
this pull request
May 4, 2026
Mosaic TPU compiler cannot reshape bool vectors (vector<Nxi1> -> vector<Nx1xi1>). When _mask_to generates combined masks like mask_q[:, None] & mask_kv[None, :], each mask is now cast to float32 before dimension expansion and combined with * instead of &. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> stack-info: PR: #2216, branch: AmesingFlank/stack/38
AmesingFlank
added a commit
that referenced
this pull request
May 4, 2026
Mosaic TPU compiler cannot reshape bool vectors (vector<Nxi1> -> vector<Nx1xi1>). When _mask_to generates combined masks like mask_q[:, None] & mask_kv[None, :], each mask is now cast to float32 before dimension expansion and combined with * instead of &. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> stack-info: PR: #2216, branch: AmesingFlank/stack/38
AmesingFlank
force-pushed
the
AmesingFlank/stack/38
branch
from
8e224a5 to
fb96c3b
Compare
AmesingFlank
added a commit
that referenced
this pull request
May 4, 2026
Mosaic TPU compiler cannot reshape bool vectors (vector<Nxi1> -> vector<Nx1xi1>). When _mask_to generates combined masks like mask_q[:, None] & mask_kv[None, :], each mask is now cast to float32 before dimension expansion and combined with * instead of &. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> stack-info: PR: #2216, branch: AmesingFlank/stack/38
AmesingFlank
force-pushed
the
AmesingFlank/stack/38
branch
from
fb96c3b to
539e8bd
Compare
AmesingFlank
added a commit
that referenced
this pull request
May 4, 2026
Mosaic TPU compiler cannot reshape bool vectors (vector<Nxi1> -> vector<Nx1xi1>). When _mask_to generates combined masks like mask_q[:, None] & mask_kv[None, :], each mask is now cast to float32 before dimension expansion and combined with * instead of &. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> stack-info: PR: #2216, branch: AmesingFlank/stack/38
AmesingFlank
force-pushed
the
AmesingFlank/stack/38
branch
from
6cdd1f0 to
332cd72
Compare
AmesingFlank
commented
May 4, 2026
| expand = state.tile_strategy.expand_str(input_sizes, dim) | ||
| expr = f"({mask_var}{expand})" | ||
| # Cast bool mask to float before expanding — Mosaic cannot | ||
| # reshape bool vectors (e.g. vector<32xi1> → vector<32x1xi1>). |
Contributor
There was a problem hiding this comment.
Can you cast it back to bool after expanding so that we can keep "&".join(mask_exprs)?
Contributor
Author
There was a problem hiding this comment.
Just gave that a try, I think casting to float, expand, and then casting back to bool, does get around this compiler error. However I'm not sure if its worth it to cast back, because we'll need to use floating point ops to mask the original float arrays anyways. (either a jnp.where as a ternary floating point op, or * as a binary floating point op)
Contributor
Author
There was a problem hiding this comment.
Tried to profile a bit, in this gist https://gist.github.com/AmesingFlank/2b3b43cc9c0ab268be0a3b14c8f8291c
not seeing any measurable difference between the two approaches. Using * is slightly less cluttering in the generated code though
Contributor
There was a problem hiding this comment.
Thanks! I think this is a similar problem as in #2214. Let's keep it for now.
AmesingFlank
added a commit
that referenced
this pull request
May 4, 2026
Mosaic TPU compiler cannot reshape bool vectors (vector<Nxi1> -> vector<Nx1xi1>). When _mask_to generates combined masks like mask_q[:, None] & mask_kv[None, :], each mask is now cast to float32 before dimension expansion and combined with * instead of &. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> stack-info: PR: #2216, branch: AmesingFlank/stack/38
AmesingFlank
added a commit
that referenced
this pull request
May 4, 2026
Mosaic TPU compiler cannot reshape bool vectors (vector<Nxi1> -> vector<Nx1xi1>). When _mask_to generates combined masks like mask_q[:, None] & mask_kv[None, :], each mask is now cast to float32 before dimension expansion and combined with * instead of &. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> stack-info: PR: #2216, branch: AmesingFlank/stack/38
AmesingFlank
force-pushed
the
AmesingFlank/stack/38
branch
from
332cd72 to
464ce56
Compare
AmesingFlank
added a commit
that referenced
this pull request
May 4, 2026
Mosaic TPU compiler cannot reshape bool vectors (vector<Nxi1> -> vector<Nx1xi1>). When _mask_to generates combined masks like mask_q[:, None] & mask_kv[None, :], each mask is now cast to float32 before dimension expansion and combined with * instead of &. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> stack-info: PR: #2216, branch: AmesingFlank/stack/38
AmesingFlank
added a commit
that referenced
this pull request
May 4, 2026
Mosaic TPU compiler cannot reshape bool vectors (vector<Nxi1> -> vector<Nx1xi1>). When _mask_to generates combined masks like mask_q[:, None] & mask_kv[None, :], each mask is now cast to float32 before dimension expansion and combined with * instead of &. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> stack-info: PR: #2216, branch: AmesingFlank/stack/38
AmesingFlank
force-pushed
the
AmesingFlank/stack/38
branch
from
464ce56 to
3d82c29
Compare
norx1991
approved these changes
May 4, 2026
jansel
approved these changes
May 5, 2026
Mosaic TPU compiler cannot reshape bool vectors (vector<Nxi1> -> vector<Nx1xi1>). When _mask_to generates combined masks like mask_q[:, None] & mask_kv[None, :], each mask is now cast to float32 before dimension expansion and combined with * instead of &. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> stack-info: PR: #2216, branch: AmesingFlank/stack/38
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked PRs:
[Pallas] Cast bool masks to float before expanding in _mask_to codegen
Mosaic TPU compiler cannot reshape bool vectors (vector ->
vector). When _mask_to generates combined masks like
mask_q[:, None] & mask_kv[None, :], each mask is now cast to float32
before dimension expansion and combined with * instead of &.
Co-Authored-By: Claude Opus 4.6 noreply@anthropic.com